diff --git a/python/src/foam_stepper/_foam_stepper.so b/python/src/foam_stepper/_foam_stepper.so index 317b645..84689a9 100755 Binary files a/python/src/foam_stepper/_foam_stepper.so and b/python/src/foam_stepper/_foam_stepper.so differ diff --git a/python/src/foam_stepper/differential_trace.py b/python/src/foam_stepper/differential_trace.py new file mode 100644 index 0000000..1ee477b --- /dev/null +++ b/python/src/foam_stepper/differential_trace.py @@ -0,0 +1,594 @@ +"""Differential trace artifacts for OpenFOAM-vs-GPU RANS parity diagnostics. + +The verifier already emits numeric split-stage artifacts. This module gives those +artifacts a stable semantic order, compact first-divergence summaries, and a +substitution contract that can materialize candidate artifacts with selected +checkpoints replaced by the OpenFOAM reference state. +""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from collections.abc import Callable, Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import numpy as np + +TRACE_SCHEMA_VERSION = 1 +DEFAULT_TOP_N = 5 +ARTIFACT_ONLY_PREFIXES = ("candidate_diagnostics.",) + +CHECKPOINT_ORDER: tuple[str, ...] = ( + "matrix_terms.UEqn.ddt.diag", + "matrix_terms.UEqn.ddt.source", + "matrix_terms.UEqn.div.diag", + "matrix_terms.UEqn.div.upper", + "matrix_terms.UEqn.div.lower", + "matrix_terms.UEqn.div.source", + "matrix_terms.UEqn.divDevSigma.diag", + "matrix_terms.UEqn.divDevSigma.upper", + "matrix_terms.UEqn.divDevSigma.lower", + "matrix_terms.UEqn.divDevSigma.wall_dev_tau.source", + "matrix_terms.UEqn.divDevSigma.source", + "matrix_operator.UEqn.unrelaxed_diag", + "matrix_operator.UEqn.unrelaxed_source", + "matrix_operator.UEqn.diag", + "matrix_operator.UEqn.upper", + "matrix_operator.UEqn.lower", + "matrix_operator.UEqn.source", + "matrix_operator.UEqn.internal_coeffs", + "matrix_operator.UEqn.boundary_coeffs", + "solver.solve_UEqn.matrix_before.diag", + "solver.solve_UEqn.matrix_before.upper", + "solver.solve_UEqn.matrix_before.lower", + "solver.solve_UEqn.matrix_before.source", + "solver.solve_UEqn.solve_diag", + "solver.solve_UEqn.solve_source", + "solver.solve_UEqn.initial_residual", + "solver.solve_UEqn.preconditioned_residual", + "solver.solve_UEqn.level_preconditioned_residual", + "solver.solve_UEqn.operator_preconditioned_direction", + "solver.solve_UEqn.first_iteration.rho", + "solver.solve_UEqn.first_iteration.denominator", + "solver.solve_UEqn.first_iteration.alpha", + "solver.solve_UEqn.first_iteration.intermediate_residual", + "solver.solve_UEqn.first_iteration.second_preconditioned_residual", + "solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual", + "solver.solve_UEqn.first_iteration.omega_numerator", + "solver.solve_UEqn.first_iteration.omega_denominator", + "solver.solve_UEqn.first_iteration.omega", + "solver.solve_UEqn.first_iteration.residual_after_omega", + "solver.solve_UEqn.first_iteration.solution_after_omega", + "solver.solve_UEqn.field_after", + "pressure_inputs.rAU", + "pressure_inputs.rAtU", + "pressure_inputs.HbyA", + "pressure_inputs.phiHbyA", + "matrix_operator.pEqn.diag", + "matrix_operator.pEqn.upper", + "matrix_operator.pEqn.lower", + "matrix_operator.pEqn.source", + "matrix_operator.pEqn.internal_coeffs", + "matrix_operator.pEqn.boundary_coeffs", + "solver.solve_pEqn.matrix_before.diag", + "solver.solve_pEqn.matrix_before.upper", + "solver.solve_pEqn.matrix_before.source", + "solver.solve_pEqn.initial_residual", + "solver.solve_pEqn.preconditioned_residual", + "solver.solve_pEqn.operator_preconditioned_direction", + "solver.solve_pEqn.p", + "solver.solve_pEqn.phi", + "final_correction.U", + "final_correction.p", + "final_correction.phi", + "turbulence.nut", + "turbulence.k", + "turbulence.omega", + "fields.U", + "fields.p", + "fields.phi", + "fields.nut", + "fields.k", + "fields.omega", +) +CHECKPOINT_ORDER_INDEX = {name: index for index, name in enumerate(CHECKPOINT_ORDER)} +REQUIRED_COMPARISON_CHECKPOINTS: frozenset[str] = frozenset( + { + "matrix_operator.UEqn.diag", + "matrix_operator.UEqn.upper", + "matrix_operator.UEqn.lower", + "matrix_operator.UEqn.source", + "matrix_operator.UEqn.psi", + "solver.solve_UEqn.matrix_before.diag", + "solver.solve_UEqn.matrix_before.upper", + "solver.solve_UEqn.matrix_before.lower", + "solver.solve_UEqn.matrix_before.source", + "solver.solve_UEqn.solve_diag", + "solver.solve_UEqn.solve_source", + "solver.solve_UEqn.initial_residual", + "solver.solve_UEqn.preconditioned_residual", + "solver.solve_UEqn.level_preconditioned_residual", + "solver.solve_UEqn.operator_preconditioned_direction", + "solver.solve_UEqn.first_iteration.rho", + "solver.solve_UEqn.first_iteration.denominator", + "solver.solve_UEqn.first_iteration.alpha", + "solver.solve_UEqn.first_iteration.intermediate_residual", + "solver.solve_UEqn.first_iteration.second_preconditioned_residual", + "solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual", + "solver.solve_UEqn.first_iteration.omega_numerator", + "solver.solve_UEqn.first_iteration.omega_denominator", + "solver.solve_UEqn.first_iteration.omega", + "solver.solve_UEqn.first_iteration.residual_after_omega", + "solver.solve_UEqn.first_iteration.solution_after_omega", + "solver.solve_UEqn.field_after", + "pressure_inputs.rAU", + "pressure_inputs.rAtU", + "pressure_inputs.HbyA", + "pressure_inputs.phiHbyA", + "matrix_operator.pEqn.diag", + "matrix_operator.pEqn.upper", + "matrix_operator.pEqn.source", + "matrix_operator.pEqn.psi", + "solver.solve_pEqn.matrix_before.diag", + "solver.solve_pEqn.matrix_before.upper", + "solver.solve_pEqn.matrix_before.source", + "solver.solve_pEqn.p", + "solver.solve_pEqn.phi", + } +) + +SUPPORTED_SUBSTITUTION_CHECKPOINTS: frozenset[str] = frozenset( + { + "matrix_operator.UEqn.diag", + "matrix_operator.UEqn.upper", + "matrix_operator.UEqn.lower", + "matrix_operator.UEqn.source", + "solver.solve_UEqn.matrix_before.diag", + "solver.solve_UEqn.matrix_before.upper", + "solver.solve_UEqn.matrix_before.lower", + "solver.solve_UEqn.matrix_before.source", + "solver.solve_UEqn.solve_source", + "solver.solve_UEqn.field_after", + "pressure_inputs.rAU", + "pressure_inputs.rAtU", + "pressure_inputs.HbyA", + "pressure_inputs.phiHbyA", + "matrix_operator.pEqn.diag", + "matrix_operator.pEqn.upper", + "matrix_operator.pEqn.source", + "solver.solve_pEqn.p", + "solver.solve_pEqn.phi", + } +) + +LocationContext = Callable[[str, tuple[int, ...], tuple[int, ...]], Mapping[str, Any] | None] + + +def finite_float(value: Any) -> float | None: + try: + number = float(value) + except (TypeError, ValueError): + return None + return number if math.isfinite(number) else None + + +def json_ready(value: Any) -> Any: + if isinstance(value, np.generic): + return value.item() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, Path): + return str(value) + if isinstance(value, Mapping): + return {str(key): json_ready(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [json_ready(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def load_npz_artifacts(path: Path | str | None) -> dict[str, np.ndarray] | None: + if path is None: + return None + artifact_path = Path(path) + if not artifact_path.exists(): + return None + with np.load(artifact_path, allow_pickle=False) as data: + return {name: np.asarray(data[name]) for name in data.files} + + +def write_json_report(path: Path | str, report: Mapping[str, Any]) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(json_ready(report), indent=2, sort_keys=True, allow_nan=False) + "\n") + + +def checkpoint_sort_key(name: str) -> tuple[int, str]: + return (CHECKPOINT_ORDER_INDEX.get(name, len(CHECKPOINT_ORDER)), name) + + +def is_trace_checkpoint(name: str) -> bool: + return not name.startswith(ARTIFACT_ONLY_PREFIXES) + + +def checkpoint_metadata(name: str) -> dict[str, Any]: + parts = name.split(".") + family = parts[0] if parts else "unknown" + metadata: dict[str, Any] = { + "checkpoint": name, + "family": family, + "order": CHECKPOINT_ORDER_INDEX.get(name), + "role": "candidate_or_reference_state", + "lifecycle_phase": lifecycle_phase(name), + } + if len(parts) >= 2: + metadata["object"] = parts[1] + if family == "matrix_operator" and len(parts) >= 3: + metadata.update({"field_or_matrix": parts[1], "coefficient": parts[2], "granularity": "cell_or_internal_face_array"}) + elif family == "matrix_terms" and len(parts) >= 5: + metadata.update({"field_or_matrix": parts[1], "term": parts[2], "subterm": parts[3], "coefficient": parts[4]}) + elif family == "matrix_terms" and len(parts) >= 4: + metadata.update({"field_or_matrix": parts[1], "term": parts[2], "coefficient": parts[3]}) + elif family == "pressure_inputs" and len(parts) >= 2: + metadata.update({"field": parts[1], "granularity": "cell_or_internal_face_array"}) + elif family == "solver" and len(parts) >= 3: + metadata.update({"solver_stage": parts[1], "solver_quantity": ".".join(parts[2:])}) + elif family in {"final_correction", "fields", "turbulence"} and len(parts) >= 2: + metadata["field"] = parts[1] + metadata["required_comparison"] = name in REQUIRED_COMPARISON_CHECKPOINTS + metadata["substitution_supported"] = name in SUPPORTED_SUBSTITUTION_CHECKPOINTS + return metadata + + +def lifecycle_phase(name: str) -> str: + if name.startswith("matrix_terms."): + return "term_contribution_before_accumulation" + if name.startswith("matrix_operator."): + if ".unrelaxed_" in name: + return "persistent_matrix_before_relaxation" + if ".internal_coeffs" in name or ".boundary_coeffs" in name: + return "matrix_boundary_coefficients" + return "persistent_matrix_after_assembly" + if ".matrix_before." in name: + return "solve_time_matrix_before_solver" + if name.startswith("pressure_inputs."): + return "pressure_input_construction" + if ".initial_residual" in name or ".preconditioned_residual" in name or ".operator_preconditioned_direction" in name: + return "linear_solver_iteration_state" + if name.startswith("solver."): + return "linear_solver_result" + if name.startswith("final_correction."): + return "pressure_velocity_flux_correction" + if name.startswith("turbulence."): + return "turbulence_update" + if name.startswith("fields."): + return "final_field_state" + return "unknown" + + +def array_stats(array: Any) -> dict[str, Any]: + arr = np.asarray(array) + out: dict[str, Any] = {"shape": [int(dim) for dim in arr.shape], "dtype": str(arr.dtype), "size": int(arr.size)} + if arr.size == 0: + return {**out, "min": None, "max": None, "mean": None, "l2": 0.0} + finite = np.isfinite(arr) + if not np.any(finite): + return {**out, "min": None, "max": None, "mean": None, "l2": None, "nonfinite_count": int(arr.size)} + finite_values = arr[finite] + out.update( + { + "min": finite_float(np.min(finite_values)), + "max": finite_float(np.max(finite_values)), + "mean": finite_float(np.mean(finite_values)), + "l2": finite_float(np.linalg.norm(finite_values.reshape(-1))), + "nonfinite_count": int(arr.size - np.count_nonzero(finite)), + } + ) + return out + + +def _value_at(array: np.ndarray, index: tuple[int, ...]) -> Any: + return json_ready(array[index]) + + +def _artifact_location(name: str, index: tuple[int, ...], shape: tuple[int, ...]) -> dict[str, Any]: + entity_kind = "array" + if name.startswith(("matrix_operator.", "pressure_inputs.rAU", "pressure_inputs.rAtU", "pressure_inputs.HbyA", "solver.solve_UEqn", "solver.solve_pEqn", "final_correction.", "fields.", "turbulence.")): + entity_kind = "cell" + if ".upper" in name or ".lower" in name or name.endswith("phi") or name.endswith("phiHbyA"): + entity_kind = "internal_face" + return { + "array_index": list(index), + "entity_kind": entity_kind, + "entity_index": int(index[0]) if index else None, + "component_index": list(index[1:]) if len(index) > 1 else None, + "artifact": name, + } + + +def _top_differences( + name: str, + candidate: np.ndarray, + reference: np.ndarray, + diff: np.ndarray, + *, + top_n: int, + location_context: LocationContext | None, +) -> list[dict[str, Any]]: + if diff.size == 0 or top_n <= 0: + return [] + flat = diff.reshape(-1) + finite = np.isfinite(flat) + if np.any(finite): + finite_indices = np.flatnonzero(finite) + finite_values = flat[finite_indices] + if finite_values.size > top_n: + local = np.argpartition(finite_values, -top_n)[-top_n:] + candidate_indices = finite_indices[local] + else: + candidate_indices = finite_indices + ordered = sorted(candidate_indices, key=lambda item: (-float(flat[int(item)]), int(item))) + else: + ordered = [int(item) for item in np.flatnonzero(~finite)[:top_n]] + + out: list[dict[str, Any]] = [] + for flat_index in ordered[:top_n]: + index = tuple(int(item) for item in np.unravel_index(int(flat_index), diff.shape)) + location = _artifact_location(name, index, tuple(diff.shape)) + context = location_context(name, index, tuple(diff.shape)) if location_context is not None else None + out.append( + { + "location": location, + "abs_difference": finite_float(diff[index]), + "candidate_value": _value_at(candidate, index), + "reference_value": _value_at(reference, index), + "local_context": json_ready(context), + } + ) + return out + + +def compare_checkpoint( + name: str, + reference: np.ndarray | None, + candidate: np.ndarray | None, + *, + rtol: float, + atol: float, + top_n: int = DEFAULT_TOP_N, + location_context: LocationContext | None = None, +) -> dict[str, Any]: + out: dict[str, Any] = { + "name": name, + "metadata": checkpoint_metadata(name), + "reference_available": reference is not None, + "candidate_available": candidate is not None, + "rtol": float(rtol), + "atol": float(atol), + } + if reference is None or candidate is None: + return {**out, "status": "missing", "allclose": False, "reason": "missing_reference" if reference is None else "missing_candidate"} + + ref = np.asarray(reference) + cand = np.asarray(candidate) + out.update( + { + "reference_stats": array_stats(ref), + "candidate_stats": array_stats(cand), + "shape_matches": ref.shape == cand.shape, + "dtype_matches": ref.dtype == cand.dtype, + } + ) + if ref.shape != cand.shape: + return {**out, "status": "failed", "allclose": False, "reason": "shape_mismatch"} + if ref.size == 0: + return {**out, "status": "passed", "allclose": True, "reason": None, "difference_stats": array_stats(np.asarray([], dtype=np.float64)), "top_differences": []} + + diff = np.abs(cand - ref) + finite = np.isfinite(diff) + finite_diff = diff[finite] + allclose = bool(np.allclose(cand, ref, rtol=rtol, atol=atol, equal_nan=False)) + if finite_diff.size: + max_flat = int(np.argmax(np.where(finite, diff, -np.inf))) + max_index = tuple(int(item) for item in np.unravel_index(max_flat, diff.shape)) + max_abs = finite_float(diff[max_index]) + mean_abs = finite_float(np.mean(finite_diff)) + rms_abs = finite_float(np.sqrt(np.mean(np.square(finite_diff)))) + else: + max_index = tuple(int(item) for item in np.unravel_index(int(np.flatnonzero(~finite.reshape(-1))[0]), diff.shape)) + max_abs = None + mean_abs = None + rms_abs = None + location = _artifact_location(name, max_index, tuple(diff.shape)) + out.update( + { + "status": "passed" if allclose else "failed", + "allclose": allclose, + "reason": None if allclose else "value_mismatch", + "difference_stats": { + "max_abs": max_abs, + "mean_abs": mean_abs, + "rms_abs": rms_abs, + "nonfinite_error_count": int(diff.size - np.count_nonzero(finite)), + }, + "largest_difference": { + "location": location, + "candidate_value": _value_at(cand, max_index), + "reference_value": _value_at(ref, max_index), + "local_context": json_ready(location_context(name, max_index, tuple(diff.shape)) if location_context is not None else None), + }, + "top_differences": _top_differences(name, cand, ref, diff, top_n=top_n, location_context=location_context), + } + ) + return out + + +def build_differential_trace_report( + *, + reference_artifacts: Mapping[str, np.ndarray] | None, + candidate_artifacts: Mapping[str, np.ndarray] | None, + reference_path: Path | str | None = None, + candidate_path: Path | str | None = None, + rtol: float = 0.0, + atol: float = 0.0, + top_n: int = DEFAULT_TOP_N, + location_context: LocationContext | None = None, + substitution: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + reference = dict(reference_artifacts or {}) + candidate = dict(candidate_artifacts or {}) + keys = sorted((name for name in set(reference) | set(candidate) if is_trace_checkpoint(name)), key=checkpoint_sort_key) + comparisons = [ + compare_checkpoint( + key, + reference.get(key), + candidate.get(key), + rtol=rtol, + atol=atol, + top_n=top_n, + location_context=location_context, + ) + for key in keys + ] + failed = [item for item in comparisons if item.get("status") == "failed"] + missing = [item for item in comparisons if item.get("status") == "missing"] + required_missing = [item for item in missing if item.get("name") in REQUIRED_COMPARISON_CHECKPOINTS] + artifacts_available = reference_artifacts is not None and candidate_artifacts is not None and bool(keys) + trace_status = "failed" if failed else "incomplete" if required_missing or not artifacts_available else "passed" + first_divergence = next((item for item in comparisons if item.get("status") == "failed"), None) + family_counts: dict[str, dict[str, int]] = defaultdict(lambda: {"passed": 0, "failed": 0, "missing": 0}) + for item in comparisons: + family = str(item.get("metadata", {}).get("family", "unknown")) + status = str(item.get("status", "missing")) + if status not in family_counts[family]: + family_counts[family][status] = 0 + family_counts[family][status] += 1 + + substitution_report = { + "schema_version": TRACE_SCHEMA_VERSION, + "supported_checkpoints": sorted(SUPPORTED_SUBSTITUTION_CHECKPOINTS, key=checkpoint_sort_key), + "replay_contract": "substitution replaces a named candidate checkpoint with the OpenFOAM reference array before downstream GPU computation; unsupported keys are artifact-only materialization targets", + **dict(substitution or {}), + } + return { + "schema_version": TRACE_SCHEMA_VERSION, + "status": trace_status, + "comparison_basis": "npz_reference_candidate_checkpoint_arrays", + "roles": { + "reference": {"path": str(reference_path) if reference_path is not None else None, "checkpoint_count": len(reference)}, + "candidate": {"path": str(candidate_path) if candidate_path is not None else None, "checkpoint_count": len(candidate)}, + }, + "rtol": float(rtol), + "atol": float(atol), + "top_n": int(top_n), + "checkpoint_count": len(keys), + "passed_count": sum(1 for item in comparisons if item.get("status") == "passed"), + "failed_count": len(failed), + "missing_count": len(missing), + "required_missing_count": len(required_missing), + "first_divergence": compact_checkpoint(first_divergence) if first_divergence is not None else None, + "first_missing": compact_checkpoint(required_missing[0] if required_missing else missing[0]) if (required_missing or missing) else None, + "families": {family: dict(counts) for family, counts in sorted(family_counts.items())}, + "checkpoints": comparisons, + "substitution": substitution_report, + } + + +def compact_checkpoint(item: Mapping[str, Any] | None) -> dict[str, Any] | None: + if item is None: + return None + diff = item.get("difference_stats", {}) if isinstance(item.get("difference_stats"), Mapping) else {} + return { + "name": item.get("name"), + "status": item.get("status"), + "reason": item.get("reason"), + "metadata": item.get("metadata"), + "max_abs": diff.get("max_abs"), + "mean_abs": diff.get("mean_abs"), + "rms_abs": diff.get("rms_abs"), + "largest_difference": item.get("largest_difference"), + "top_differences": item.get("top_differences", []), + } + + +def materialize_substituted_artifacts( + *, + reference_artifacts: Mapping[str, np.ndarray], + candidate_artifacts: Mapping[str, np.ndarray], + checkpoints: Sequence[str], + output_path: Path | str, +) -> dict[str, Any]: + requested = list(dict.fromkeys(str(item) for item in checkpoints)) + substituted = {name: np.ascontiguousarray(value) for name, value in candidate_artifacts.items()} + applied: list[dict[str, Any]] = [] + missing_reference: list[str] = [] + missing_candidate: list[str] = [] + shape_mismatches: list[dict[str, Any]] = [] + + for checkpoint in requested: + ref = reference_artifacts.get(checkpoint) + cand = candidate_artifacts.get(checkpoint) + if ref is None: + missing_reference.append(checkpoint) + continue + if cand is None: + missing_candidate.append(checkpoint) + continue + ref_array = np.asarray(ref) + cand_array = np.asarray(cand) + if ref_array.shape != cand_array.shape: + shape_mismatches.append({"checkpoint": checkpoint, "reference_shape": list(ref_array.shape), "candidate_shape": list(cand_array.shape)}) + continue + substituted[checkpoint] = np.ascontiguousarray(ref_array) + applied.append( + { + "checkpoint": checkpoint, + "shape": [int(dim) for dim in ref_array.shape], + "dtype": str(ref_array.dtype), + "gpu_replay_supported": checkpoint in SUPPORTED_SUBSTITUTION_CHECKPOINTS, + } + ) + + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(output, **substituted) + status = "ready" if applied and not missing_reference and not missing_candidate and not shape_mismatches else "partial" if applied else "failed" + return { + "schema_version": TRACE_SCHEMA_VERSION, + "status": status, + "path": str(output), + "format": "npz", + "requested": requested, + "applied": applied, + "missing_reference": missing_reference, + "missing_candidate": missing_candidate, + "shape_mismatches": shape_mismatches, + "artifact_count": len(substituted), + } + + +def extract_reference_substitutions(reference_artifacts: Mapping[str, np.ndarray], checkpoints: Iterable[str]) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + substitutions: dict[str, np.ndarray] = {} + missing: list[str] = [] + unsupported: list[str] = [] + for checkpoint in dict.fromkeys(str(item) for item in checkpoints): + if checkpoint not in SUPPORTED_SUBSTITUTION_CHECKPOINTS: + unsupported.append(checkpoint) + continue + value = reference_artifacts.get(checkpoint) + if value is None: + missing.append(checkpoint) + continue + substitutions[checkpoint] = np.ascontiguousarray(np.asarray(value, dtype=np.float64)) + return substitutions, { + "schema_version": TRACE_SCHEMA_VERSION, + "requested": list(dict.fromkeys(str(item) for item in checkpoints)), + "loaded": sorted(substitutions, key=checkpoint_sort_key), + "missing_reference": missing, + "unsupported": unsupported, + "status": "ready" if substitutions and not missing and not unsupported else "partial" if substitutions else "failed" if checkpoints else "not_requested", + } diff --git a/python/src/foam_stepper/gpu/backend.py b/python/src/foam_stepper/gpu/backend.py index 813e7ad..1bfe17b 100644 --- a/python/src/foam_stepper/gpu/backend.py +++ b/python/src/foam_stepper/gpu/backend.py @@ -19,6 +19,7 @@ from typing import Any import numpy as np import quadrants as qd +from ..differential_trace import SUPPORTED_SUBSTITUTION_CHECKPOINTS from .constants import ( DEFAULT_LAMINAR_NU, @@ -68,6 +69,7 @@ from .kernels import ( gpu_rans_momentum_gauss_grad_u_internal, gpu_rans_momentum_gauss_grad_u_boundary, gpu_rans_momentum_gauss_grad_u_finish, + gpu_rans_momentum_gauss_grad_u_by_cell, gpu_rans_momentum_internal_dev_tau_source, gpu_rans_momentum_wall_dev_tau_source, gpu_rans_momentum_linear_upwind_source, @@ -107,11 +109,17 @@ from .linear_solve import ( gpu_bicgstab_dot_vector, gpu_bicgstab_precondition_vector, gpu_copy_scalar, + gpu_cg_dot_scalar, + gpu_dic_apply_scalar_symmetric_levels, gpu_copy_vector, gpu_dilu_apply_vector_asymmetric_faces, + gpu_dilu_apply_vector_asymmetric_levels, gpu_ldu_matvec_vector_asymmetric_faces, gpu_ldu_pbicgstab_vector_asymmetric_components, + gpu_ldu_pbicgstab_vector_asymmetric_first_iteration_components, + gpu_ldu_matvec_scalar_symmetric_faces, gpu_ldu_pcg_scalar_symmetric_faces, + gpu_pcg_initialize_residual_scalar, gpu_zero_scalar_accumulator, ) @@ -1020,6 +1028,77 @@ def build_openfoam_offdiag_cell_faces( ) +def build_openfoam_face_order_cell_faces( + n_cells: int, + owner: np.ndarray, + neighbour: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + owner_i32 = np.asarray(owner, dtype=np.int32).reshape(-1) + neighbour_i32 = np.asarray(neighbour, dtype=np.int32).reshape(-1) + if owner_i32.shape != neighbour_i32.shape: + raise ValueError(f"owner/neighbour shape mismatch: {owner_i32.shape} != {neighbour_i32.shape}") + if owner_i32.size == 0: + return np.zeros(n_cells + 1, dtype=np.int32), np.zeros(0, dtype=np.int32), np.zeros(0, dtype=np.int32) + if int(owner_i32.min()) < 0 or int(neighbour_i32.min()) < 0 or int(max(owner_i32.max(), neighbour_i32.max())) >= n_cells: + raise ValueError("owner/neighbour addresses exceed cell range") + + counts = np.zeros(n_cells, dtype=np.int32) + np.add.at(counts, owner_i32, 1) + np.add.at(counts, neighbour_i32, 1) + offsets = np.empty(n_cells + 1, dtype=np.int32) + offsets[0] = 0 + np.cumsum(counts, out=offsets[1:]) + faces = np.empty(int(offsets[-1]), dtype=np.int32) + sides = np.empty(int(offsets[-1]), dtype=np.int32) + cursor = offsets[:-1].copy() + + for face, (owner_cell, neighbour_cell) in enumerate(zip(owner_i32, neighbour_i32, strict=True)): + owner_slot = int(cursor[int(owner_cell)]) + faces[owner_slot] = int(face) + sides[owner_slot] = 0 + cursor[int(owner_cell)] += 1 + + neighbour_slot = int(cursor[int(neighbour_cell)]) + faces[neighbour_slot] = int(face) + sides[neighbour_slot] = 1 + cursor[int(neighbour_cell)] += 1 + + return ( + np.ascontiguousarray(offsets, dtype=np.int32), + np.ascontiguousarray(faces, dtype=np.int32), + np.ascontiguousarray(sides, dtype=np.int32), + ) + + +def build_boundary_cell_faces( + n_cells: int, + face_cells: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + face_cells_i32 = np.asarray(face_cells, dtype=np.int32).reshape(-1) + if face_cells_i32.size == 0: + return np.zeros(n_cells + 1, dtype=np.int32), np.zeros(0, dtype=np.int32) + if int(face_cells_i32.min()) < 0 or int(face_cells_i32.max()) >= n_cells: + raise ValueError("boundary face-cell addresses exceed cell range") + + counts = np.zeros(n_cells, dtype=np.int32) + np.add.at(counts, face_cells_i32, 1) + offsets = np.empty(n_cells + 1, dtype=np.int32) + offsets[0] = 0 + np.cumsum(counts, out=offsets[1:]) + faces = np.empty(int(offsets[-1]), dtype=np.int32) + cursor = offsets[:-1].copy() + + for face, cell in enumerate(face_cells_i32): + slot = int(cursor[int(cell)]) + faces[slot] = int(face) + cursor[int(cell)] += 1 + + return ( + np.ascontiguousarray(offsets, dtype=np.int32), + np.ascontiguousarray(faces, dtype=np.int32), + ) + + def openfoam_dilu_preconditioner_diag( diag: np.ndarray, owner: np.ndarray, @@ -1145,6 +1224,60 @@ def gpu_empty_f64(shape: tuple[int, ...], path: str) -> tuple[Any, dict[str, Any gpu_array = qd.ndarray(qd.f64, shape=shape) return gpu_array, {"path": path, "gpu_shape": list(shape), "gpu_dtype": "f64", "allocated": True} +def normalize_trace_substitutions(trace_substitutions: Mapping[str, Any] | None) -> dict[str, np.ndarray]: + if not trace_substitutions: + return {} + normalized: dict[str, np.ndarray] = {} + for key, value in trace_substitutions.items(): + checkpoint = str(key) + if checkpoint not in SUPPORTED_SUBSTITUTION_CHECKPOINTS: + raise BackendExecutionError( + "prepare_trace_substitutions", + f"unsupported differential trace substitution checkpoint {checkpoint!r}", + details={ + "checkpoint": checkpoint, + "supported_checkpoints": sorted(SUPPORTED_SUBSTITUTION_CHECKPOINTS), + "used_cpu_fallback": False, + }, + ) + normalized[checkpoint] = np.ascontiguousarray(np.asarray(value, dtype=np.float64)) + return normalized + + +def apply_trace_substitution( + checkpoint: str, + gpu_array: Any, + expected: np.ndarray, + substitutions: Mapping[str, np.ndarray], + applied: list[dict[str, Any]], +) -> np.ndarray: + replacement = substitutions.get(checkpoint) + if replacement is None: + return expected + expected_shape = tuple(np.asarray(expected).shape) + if tuple(replacement.shape) != expected_shape: + raise BackendExecutionError( + "apply_trace_substitution", + f"{checkpoint}: reference substitution shape does not match candidate checkpoint", + details={ + "checkpoint": checkpoint, + "reference_shape": array_shape(replacement), + "candidate_shape": array_shape(expected), + "used_cpu_fallback": False, + }, + ) + gpu_array.from_numpy(replacement) + applied.append( + { + "checkpoint": checkpoint, + "shape": array_shape(replacement), + "dtype": str(replacement.dtype), + "role": "substituted_reference_state", + "used_cpu_fallback": False, + } + ) + return replacement + def read_case_laminar_nu(case: Path) -> float: for relative in ("constant/physicalProperties", "constant/transportProperties", "constant/transportProperties.v2112"): @@ -1384,7 +1517,9 @@ def quadrants_kernel_evidence(expected: Iterable[str]) -> dict[str, Any]: } -def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None) -> dict[str, Any]: +def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None, trace_substitutions: Mapping[str, Any] | None = None) -> dict[str, Any]: + trace_substitution_arrays = normalize_trace_substitutions(trace_substitutions) + applied_trace_substitutions: list[dict[str, Any]] = [] fields = read_fields(stepper, "gpu_stage_smoke") gpu_solver_started = time.perf_counter() input_transfer_started = time.perf_counter() @@ -1406,6 +1541,10 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P offdiag_cell_offsets_gpu, offdiag_cell_offsets_np, offdiag_cell_offsets_transfer = gpu_i32_array(offdiag_cell_offsets_np, "mesh.connectivity.offdiag_cell_offsets.internal") offdiag_cell_faces_gpu, offdiag_cell_faces_np, offdiag_cell_faces_transfer = gpu_i32_array(offdiag_cell_faces_np, "mesh.connectivity.offdiag_cell_faces.internal") offdiag_cell_sides_gpu, offdiag_cell_sides_np, offdiag_cell_sides_transfer = gpu_i32_array(offdiag_cell_sides_np, "mesh.connectivity.offdiag_cell_sides.internal") + grad_cell_offsets_np, grad_cell_faces_np, grad_cell_sides_np = build_openfoam_face_order_cell_faces(n_cells, owner_np, neighbour_np) + grad_cell_offsets_gpu, grad_cell_offsets_np, grad_cell_offsets_transfer = gpu_i32_array(grad_cell_offsets_np, "mesh.connectivity.grad_cell_offsets.openfoam_face_order") + grad_cell_faces_gpu, grad_cell_faces_np, grad_cell_faces_transfer = gpu_i32_array(grad_cell_faces_np, "mesh.connectivity.grad_cell_faces.openfoam_face_order") + grad_cell_sides_gpu, grad_cell_sides_np, grad_cell_sides_transfer = gpu_i32_array(grad_cell_sides_np, "mesh.connectivity.grad_cell_sides.openfoam_face_order") u_dilu_level_schedule = build_ldu_level_schedule(n_cells, owner_np, neighbour_np, name="solve_UEqn.DILU") p_dic_level_schedule = build_ldu_level_schedule(n_cells, owner_np, neighbour_np, name="solve_pEqn.DIC") sf_gpu, sf_np, sf_transfer = gpu_f64_array(np.asarray(mesh.Sf)[:n_internal_faces], "mesh.geometry.Sf.internal") @@ -1535,6 +1674,9 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P momentum_u_boundary_face_cells_gpu, momentum_u_boundary_face_cells_np, momentum_u_boundary_face_cells_transfer = gpu_i32_array(momentum_u_boundary_face_cells_np, "mesh.boundary.face_cells.momentum_grad_u") momentum_u_boundary_values_gpu, momentum_u_boundary_values_np, momentum_u_boundary_values_transfer = gpu_f64_array(momentum_u_boundary_values_np, "fields.U.boundary.momentum_grad_u") momentum_u_boundary_sf_gpu, momentum_u_boundary_sf_np, momentum_u_boundary_sf_transfer = gpu_f64_array(momentum_u_boundary_sf_np, "mesh.boundary.Sf.momentum_grad_u") + momentum_u_boundary_cell_offsets_np, momentum_u_boundary_cell_faces_np = build_boundary_cell_faces(n_cells, momentum_u_boundary_face_cells_np) + momentum_u_boundary_cell_offsets_gpu, momentum_u_boundary_cell_offsets_np, momentum_u_boundary_cell_offsets_transfer = gpu_i32_array(momentum_u_boundary_cell_offsets_np, "mesh.boundary.grad_u_cell_offsets.openfoam_patch_order") + momentum_u_boundary_cell_faces_gpu, momentum_u_boundary_cell_faces_np, momentum_u_boundary_cell_faces_transfer = gpu_i32_array(momentum_u_boundary_cell_faces_np, "mesh.boundary.grad_u_cell_faces.openfoam_patch_order") momentum_u_convection_face_cells_parts: list[np.ndarray] = [] momentum_u_convection_phi_parts: list[np.ndarray] = [] momentum_u_convection_internal_coeff_parts: list[np.ndarray] = [] @@ -1778,6 +1920,28 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P momentum_wall_face_centres_np = np.concatenate(momentum_wall_face_centres_parts) if momentum_wall_face_centres_parts else np.empty((0, 3), dtype=np.float64) momentum_wall_area_vectors_np = np.concatenate(momentum_wall_area_vectors_parts) if momentum_wall_area_vectors_parts else np.empty((0, 3), dtype=np.float64) momentum_wall_area_magnitudes_np = np.concatenate(momentum_wall_area_magnitudes_parts) if momentum_wall_area_magnitudes_parts else np.empty((0,), dtype=np.float64) + momentum_wall_normals_np = np.divide( + momentum_wall_area_vectors_np, + momentum_wall_area_magnitudes_np[:, None], + out=np.zeros_like(momentum_wall_area_vectors_np), + where=momentum_wall_area_magnitudes_np[:, None] > 0.0, + ) + momentum_wall_delta_np = momentum_wall_face_centres_np - cell_centres_np[momentum_wall_face_cells_np] + momentum_wall_projected_delta_np = np.abs(np.sum(momentum_wall_delta_np * momentum_wall_normals_np, axis=1)) + momentum_wall_delta_coeffs_np = np.ascontiguousarray( + 1.0 / np.maximum(momentum_wall_projected_delta_np, 1.0e-300), + dtype=np.float64, + ) + momentum_wall_gamma_np = laminar_nu + momentum_wall_nut_np + momentum_wall_gamma_area_vectors_np = np.ascontiguousarray( + momentum_wall_gamma_np[:, None] * momentum_wall_area_vectors_np, + dtype=np.float64, + ) + momentum_wall_trace_gamma_area_vectors_np = np.ascontiguousarray( + (2.0 / 3.0) * momentum_wall_gamma_area_vectors_np, + dtype=np.float64, + ) + momentum_wall_normals_np = np.ascontiguousarray(momentum_wall_normals_np, dtype=np.float64) n_momentum_wall_faces = int(momentum_wall_face_cells_np.shape[0]) momentum_wall_face_cells_gpu, momentum_wall_face_cells_np, momentum_wall_face_cells_transfer = gpu_i32_array(momentum_wall_face_cells_np, "mesh.boundary.face_cells.momentum_wall_diffusion") momentum_wall_values_gpu, momentum_wall_values_np, momentum_wall_values_transfer = gpu_f64_array(momentum_wall_values_np, "fields.U.boundary.momentum_wall_diffusion") @@ -1785,6 +1949,10 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P momentum_wall_face_centres_gpu, momentum_wall_face_centres_np, momentum_wall_face_centres_transfer = gpu_f64_array(momentum_wall_face_centres_np, "mesh.boundary.Cf.momentum_wall_diffusion") momentum_wall_area_vectors_gpu, momentum_wall_area_vectors_np, momentum_wall_area_vectors_transfer = gpu_f64_array(momentum_wall_area_vectors_np, "mesh.boundary.Sf.momentum_wall_diffusion") momentum_wall_area_magnitudes_gpu, momentum_wall_area_magnitudes_np, momentum_wall_area_magnitudes_transfer = gpu_f64_array(momentum_wall_area_magnitudes_np, "mesh.boundary.magSf.momentum_wall_diffusion") + momentum_wall_delta_coeffs_gpu, momentum_wall_delta_coeffs_np, momentum_wall_delta_coeffs_transfer = gpu_f64_array(momentum_wall_delta_coeffs_np, "mesh.boundary.deltaCoeffs.momentum_wall_dev_tau") + momentum_wall_normals_gpu, momentum_wall_normals_np, momentum_wall_normals_transfer = gpu_f64_array(momentum_wall_normals_np, "mesh.boundary.nf.momentum_wall_dev_tau") + momentum_wall_gamma_area_vectors_gpu, momentum_wall_gamma_area_vectors_np, momentum_wall_gamma_area_vectors_transfer = gpu_f64_array(momentum_wall_gamma_area_vectors_np, "mesh.boundary.gammaSf.momentum_wall_dev_tau") + momentum_wall_trace_gamma_area_vectors_gpu, momentum_wall_trace_gamma_area_vectors_np, momentum_wall_trace_gamma_area_vectors_transfer = gpu_f64_array(momentum_wall_trace_gamma_area_vectors_np, "mesh.boundary.twoThirdsGammaSf.momentum_wall_dev_tau") omega_wall_distances_by_cell = np.full((n_cells,), np.inf, dtype=np.float64) omega_wall_seed_count = 0 for patch in mesh.boundary: @@ -1846,6 +2014,15 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_upper_gpu, u_upper_alloc = gpu_empty_f64((n_internal_faces,), "gpu_stages.UEqn.upper") u_lower_gpu, u_lower_alloc = gpu_empty_f64((n_internal_faces,), "gpu_stages.UEqn.lower") grad_u_gpu, grad_u_alloc = gpu_empty_f64((n_cells, 3, 3), "gpu_stages.UEqn.gradU") + u_term_div_diag_gpu, u_term_div_diag_alloc = gpu_empty_f64((n_cells,), "gpu_stages.UEqn.terms.div.diag_scratch") + u_term_div_source_gpu, u_term_div_source_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.UEqn.terms.div.source") + u_term_stress_boundary_relax_add_gpu, u_term_stress_boundary_relax_add_alloc = gpu_empty_f64((n_cells,), "gpu_stages.UEqn.terms.divDevSigma.boundary_relax_add_scratch") + u_term_stress_boundary_relax_subtract_gpu, u_term_stress_boundary_relax_subtract_alloc = gpu_empty_f64((n_cells,), "gpu_stages.UEqn.terms.divDevSigma.boundary_relax_subtract_scratch") + u_term_stress_boundary_diag_gpu, u_term_stress_boundary_diag_alloc = gpu_empty_f64((n_cells,), "gpu_stages.UEqn.terms.divDevSigma.boundary_diag_scratch") + u_term_stress_source_gpu, u_term_stress_source_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.UEqn.terms.divDevSigma.source") + u_term_stress_wall_diffusion_source_gpu, u_term_stress_wall_diffusion_source_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.UEqn.terms.divDevSigma.wall_diffusion.source") + u_term_stress_internal_dev_tau_source_gpu, u_term_stress_internal_dev_tau_source_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.UEqn.terms.divDevSigma.internal_dev_tau.source") + u_term_stress_wall_dev_tau_source_gpu, u_term_stress_wall_dev_tau_source_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.UEqn.terms.divDevSigma.wall_dev_tau.source") rAU_gpu, rAU_alloc = gpu_empty_f64((n_cells,), "gpu_stages.pressure_inputs.rAU") H1_gpu, H1_alloc = gpu_empty_f64((n_cells,), "gpu_stages.UEqn.H1") rAtU_gpu, rAtU_alloc = gpu_empty_f64((n_cells,), "gpu_stages.pressure_inputs.rAtU") @@ -1864,6 +2041,12 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_operator_direction_gpu, u_operator_direction_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.operator_direction") u_intermediate_gpu, u_intermediate_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.intermediate") u_operator_intermediate_gpu, u_operator_intermediate_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.operator_intermediate") + u_level_preconditioned_gpu, u_level_preconditioned_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.level_preconditioned_residual") + u_first_intermediate_residual_gpu, u_first_intermediate_residual_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.first_iteration.intermediate_residual") + u_first_second_preconditioned_gpu, u_first_second_preconditioned_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.first_iteration.second_preconditioned_residual") + u_first_operator_second_preconditioned_gpu, u_first_operator_second_preconditioned_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.first_iteration.operator_second_preconditioned_residual") + u_first_residual_after_omega_gpu, u_first_residual_after_omega_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.first_iteration.residual_after_omega") + u_first_solution_after_omega_gpu, u_first_solution_after_omega_alloc = gpu_empty_f64(tuple(u_np.shape), "gpu_stages.solve_UEqn.first_iteration.solution_after_omega") u_rr_gpu, u_rr_alloc = gpu_empty_f64((1,), "gpu_stages.solve_UEqn.residual_squared") u_rho_gpu, u_rho_alloc = gpu_empty_f64((1,), "gpu_stages.solve_UEqn.rho") u_denominator_gpu, u_denominator_alloc = gpu_empty_f64((1,), "gpu_stages.solve_UEqn.denominator") @@ -1903,31 +2086,94 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P gpu_rans_momentum_bounded_convection_sp_boundary(n_pressure_boundary_faces, boundary_face_cells_gpu, boundary_phi_gpu, u_diag_gpu) if n_momentum_u_convection_boundary_faces: gpu_rans_momentum_convection_boundary_coefficients(n_momentum_u_convection_boundary_faces, momentum_u_convection_face_cells_gpu, momentum_u_convection_phi_gpu, momentum_u_convection_internal_coeff_gpu, momentum_u_convection_boundary_coeff_gpu, u_diag_gpu, u_source_gpu) - gpu_rans_zero_tensor_field(n_cells, grad_u_gpu) - gpu_rans_momentum_gauss_grad_u_internal(n_internal_faces, owner_gpu, neighbour_gpu, u_gpu, sf_gpu, face_weights_gpu, grad_u_gpu) - if n_momentum_u_boundary_faces: - gpu_rans_momentum_gauss_grad_u_boundary(n_momentum_u_boundary_faces, momentum_u_boundary_face_cells_gpu, momentum_u_boundary_values_gpu, momentum_u_boundary_sf_gpu, grad_u_gpu) - gpu_rans_momentum_gauss_grad_u_finish(n_cells, cell_volumes_gpu, grad_u_gpu) + gpu_rans_momentum_gauss_grad_u_by_cell(n_cells, grad_cell_offsets_gpu, grad_cell_faces_gpu, grad_cell_sides_gpu, owner_gpu, neighbour_gpu, u_gpu, sf_gpu, face_weights_gpu, momentum_u_boundary_cell_offsets_gpu, momentum_u_boundary_cell_faces_gpu, momentum_u_boundary_values_gpu, momentum_u_boundary_sf_gpu, cell_volumes_gpu, grad_u_gpu) gpu_rans_momentum_internal_dev_tau_source(n_internal_faces, owner_gpu, neighbour_gpu, nut_gpu, laminar_nu, cell_centres_gpu, sf_gpu, mag_sf_gpu, face_weights_gpu, grad_u_gpu, u_source_gpu) if n_momentum_wall_faces: - gpu_rans_momentum_wall_dev_tau_source(n_momentum_wall_faces, momentum_wall_face_cells_gpu, momentum_wall_values_gpu, momentum_wall_nut_gpu, laminar_nu, u_gpu, cell_centres_gpu, momentum_wall_face_centres_gpu, momentum_wall_area_vectors_gpu, momentum_wall_area_magnitudes_gpu, grad_u_gpu, u_source_gpu) + gpu_rans_momentum_wall_dev_tau_source( + n_momentum_wall_faces, + momentum_wall_face_cells_gpu, + momentum_wall_values_gpu, + momentum_wall_gamma_area_vectors_gpu, + momentum_wall_trace_gamma_area_vectors_gpu, + momentum_wall_delta_coeffs_gpu, + momentum_wall_normals_gpu, + u_gpu, + grad_u_gpu, + u_source_gpu, + ) gpu_rans_momentum_linear_upwind_source(n_internal_faces, owner_gpu, neighbour_gpu, phi_gpu, cell_centres_gpu, face_centres_gpu, grad_u_gpu, u_source_gpu) + gpu_rans_momentum_assembly(n_cells, u_gpu, cell_volumes_gpu, u_term_div_diag_gpu, u_term_div_source_gpu) + if n_momentum_u_convection_boundary_faces: + gpu_rans_momentum_convection_boundary_coefficients(n_momentum_u_convection_boundary_faces, momentum_u_convection_face_cells_gpu, momentum_u_convection_phi_gpu, momentum_u_convection_internal_coeff_gpu, momentum_u_convection_boundary_coeff_gpu, u_term_div_diag_gpu, u_term_div_source_gpu) + gpu_rans_momentum_linear_upwind_source(n_internal_faces, owner_gpu, neighbour_gpu, phi_gpu, cell_centres_gpu, face_centres_gpu, grad_u_gpu, u_term_div_source_gpu) + gpu_rans_momentum_assembly(n_cells, u_gpu, cell_volumes_gpu, u_term_stress_boundary_diag_gpu, u_term_stress_source_gpu) + gpu_rans_zero_scalar_field(n_cells, u_term_stress_boundary_relax_add_gpu) + gpu_rans_zero_scalar_field(n_cells, u_term_stress_boundary_relax_subtract_gpu) + if n_momentum_wall_faces: + gpu_rans_momentum_wall_diffusion_coefficients(n_momentum_wall_faces, momentum_wall_face_cells_gpu, momentum_wall_values_gpu, momentum_wall_nut_gpu, laminar_nu, cell_centres_gpu, momentum_wall_face_centres_gpu, momentum_wall_area_vectors_gpu, momentum_wall_area_magnitudes_gpu, u_term_stress_boundary_relax_add_gpu, u_term_stress_boundary_relax_subtract_gpu, u_term_stress_boundary_diag_gpu, u_term_stress_source_gpu) + gpu_rans_momentum_internal_dev_tau_source(n_internal_faces, owner_gpu, neighbour_gpu, nut_gpu, laminar_nu, cell_centres_gpu, sf_gpu, mag_sf_gpu, face_weights_gpu, grad_u_gpu, u_term_stress_source_gpu) + if n_momentum_wall_faces: + gpu_rans_momentum_wall_dev_tau_source( + n_momentum_wall_faces, + momentum_wall_face_cells_gpu, + momentum_wall_values_gpu, + momentum_wall_gamma_area_vectors_gpu, + momentum_wall_trace_gamma_area_vectors_gpu, + momentum_wall_delta_coeffs_gpu, + momentum_wall_normals_gpu, + u_gpu, + grad_u_gpu, + u_term_stress_source_gpu, + ) + gpu_rans_momentum_assembly(n_cells, u_gpu, cell_volumes_gpu, u_term_stress_boundary_diag_gpu, u_term_stress_wall_diffusion_source_gpu) + gpu_rans_zero_scalar_field(n_cells, u_term_stress_boundary_relax_add_gpu) + gpu_rans_zero_scalar_field(n_cells, u_term_stress_boundary_relax_subtract_gpu) + if n_momentum_wall_faces: + gpu_rans_momentum_wall_diffusion_coefficients(n_momentum_wall_faces, momentum_wall_face_cells_gpu, momentum_wall_values_gpu, momentum_wall_nut_gpu, laminar_nu, cell_centres_gpu, momentum_wall_face_centres_gpu, momentum_wall_area_vectors_gpu, momentum_wall_area_magnitudes_gpu, u_term_stress_boundary_relax_add_gpu, u_term_stress_boundary_relax_subtract_gpu, u_term_stress_boundary_diag_gpu, u_term_stress_wall_diffusion_source_gpu) + gpu_rans_momentum_assembly(n_cells, u_gpu, cell_volumes_gpu, u_term_div_diag_gpu, u_term_stress_internal_dev_tau_source_gpu) + gpu_rans_momentum_internal_dev_tau_source(n_internal_faces, owner_gpu, neighbour_gpu, nut_gpu, laminar_nu, cell_centres_gpu, sf_gpu, mag_sf_gpu, face_weights_gpu, grad_u_gpu, u_term_stress_internal_dev_tau_source_gpu) + gpu_rans_momentum_assembly(n_cells, u_gpu, cell_volumes_gpu, u_term_div_diag_gpu, u_term_stress_wall_dev_tau_source_gpu) + if n_momentum_wall_faces: + gpu_rans_momentum_wall_dev_tau_source( + n_momentum_wall_faces, + momentum_wall_face_cells_gpu, + momentum_wall_values_gpu, + momentum_wall_gamma_area_vectors_gpu, + momentum_wall_trace_gamma_area_vectors_gpu, + momentum_wall_delta_coeffs_gpu, + momentum_wall_normals_gpu, + u_gpu, + grad_u_gpu, + u_term_stress_wall_dev_tau_source_gpu, + ) gpu_copy_scalar(n_cells, u_diag_gpu, u_unrelaxed_diag_gpu) gpu_copy_vector(n_cells, u_source_gpu, u_unrelaxed_source_gpu) gpu_copy_scalar(n_cells, momentum_internal_diag_gpu, H1_gpu) if n_momentum_u_convection_boundary_diag_faces: gpu_rans_momentum_boundary_relaxation_coefficients(n_momentum_u_convection_boundary_diag_faces, momentum_u_convection_boundary_diag_face_cells_gpu, momentum_u_convection_boundary_diag_coeff_gpu, u_boundary_relax_add_gpu, u_boundary_relax_subtract_gpu, u_boundary_diag_coeff_gpu) gpu_rans_momentum_equation_relaxation(n_cells, u_gpu, DEFAULT_MOMENTUM_RELAXATION_ALPHA, H1_gpu, u_boundary_relax_add_gpu, u_boundary_relax_subtract_gpu, u_diag_gpu, u_source_gpu) + if trace_substitution_arrays: + apply_trace_substitution("matrix_operator.UEqn.diag", u_diag_gpu, np.asarray(u_diag_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("matrix_operator.UEqn.upper", u_upper_gpu, np.asarray(u_upper_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("matrix_operator.UEqn.lower", u_lower_gpu, np.asarray(u_lower_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("matrix_operator.UEqn.source", u_source_gpu, np.asarray(u_source_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_copy_scalar(n_cells, u_diag_gpu, u_boundary_diag_candidate_gpu) gpu_rans_add_scalar_field(n_cells, u_boundary_diag_coeff_gpu, u_boundary_diag_candidate_gpu) gpu_copy_vector(n_cells, u_source_gpu, u_matrix_source_gpu) if n_momentum_u_convection_boundary_source_faces: gpu_rans_momentum_convection_boundary_source(n_momentum_u_convection_boundary_source_faces, momentum_u_convection_boundary_source_face_cells_gpu, momentum_u_convection_boundary_source_gpu, u_source_gpu) gpu_rans_add_vector_field(n_cells, momentum_pressure_source_gpu, u_source_gpu) - u_boundary_diag_candidate_for_preconditioner = np.asarray(u_boundary_diag_candidate_gpu.to_numpy(), dtype=np.float64) - u_upper_for_preconditioner = np.asarray(u_upper_gpu.to_numpy(), dtype=np.float64) - u_lower_for_preconditioner = np.asarray(u_lower_gpu.to_numpy(), dtype=np.float64) - u_source_for_solver = np.asarray(u_source_gpu.to_numpy(), dtype=np.float64) + if trace_substitution_arrays: + u_boundary_diag_candidate_for_preconditioner = apply_trace_substitution("solver.solve_UEqn.matrix_before.diag", u_boundary_diag_candidate_gpu, np.asarray(u_boundary_diag_candidate_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + u_upper_for_preconditioner = apply_trace_substitution("solver.solve_UEqn.matrix_before.upper", u_upper_gpu, np.asarray(u_upper_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + u_lower_for_preconditioner = apply_trace_substitution("solver.solve_UEqn.matrix_before.lower", u_lower_gpu, np.asarray(u_lower_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + u_source_for_solver = apply_trace_substitution("solver.solve_UEqn.matrix_before.source", u_source_gpu, np.asarray(u_source_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + u_source_for_solver = apply_trace_substitution("solver.solve_UEqn.solve_source", u_source_gpu, u_source_for_solver, trace_substitution_arrays, applied_trace_substitutions) + else: + u_boundary_diag_candidate_for_preconditioner = np.asarray(u_boundary_diag_candidate_gpu.to_numpy(), dtype=np.float64) + u_upper_for_preconditioner = np.asarray(u_upper_gpu.to_numpy(), dtype=np.float64) + u_lower_for_preconditioner = np.asarray(u_lower_gpu.to_numpy(), dtype=np.float64) + u_source_for_solver = np.asarray(u_source_gpu.to_numpy(), dtype=np.float64) u_normalization_factors = np.asarray( [ openfoam_ldu_normalization_factor( @@ -2002,11 +2248,22 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_residual_gpu, u_operator_direction_gpu, ) + gpu_dilu_apply_vector_asymmetric_levels( + u_dilu_level_schedule, + owner_gpu, + neighbour_gpu, + u_upper_gpu, + u_lower_gpu, + u_dilu_reciprocal_diag_gpu, + u_residual_gpu, + u_level_preconditioned_gpu, + ) qd.sync() u_initial_residual = np.asarray(u_residual_gpu.to_numpy(), dtype=np.float64) u_initial_residual_squared = float(np.asarray(u_rr_gpu.to_numpy())[0]) u_diagonal_preconditioned = np.asarray(u_intermediate_gpu.to_numpy(), dtype=np.float64) u_gpu_dilu_preconditioned = np.asarray(u_operator_direction_gpu.to_numpy(), dtype=np.float64) + u_gpu_level_preconditioned = np.asarray(u_level_preconditioned_gpu.to_numpy(), dtype=np.float64) u_openfoam_dilu_preconditioned = openfoam_dilu_apply_vector_reference( owner_np, neighbour_np, @@ -2033,6 +2290,35 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_operator_preconditioned_direction = np.asarray(u_intermediate_gpu.to_numpy(), dtype=np.float64) u_first_denominator = float(np.asarray(u_denominator_gpu.to_numpy())[0]) u_first_alpha = u_initial_residual_squared / u_first_denominator if abs(u_first_denominator) > 1.0e-300 else None + u_first_iteration_diagnostics = gpu_ldu_pbicgstab_vector_asymmetric_first_iteration_components( + n_cells, + n_internal_faces, + owner_gpu, + neighbour_gpu, + u_upper_gpu, + u_lower_gpu, + u_boundary_diag_candidate_gpu, + u_source_gpu, + u_gpu, + u_first_intermediate_residual_gpu, + u_first_second_preconditioned_gpu, + u_first_operator_second_preconditioned_gpu, + u_first_residual_after_omega_gpu, + u_first_solution_after_omega_gpu, + u_rr_gpu, + u_rho_gpu, + u_denominator_gpu, + u_omega_numerator_gpu, + u_omega_denominator_gpu, + preconditioner_diag=u_preconditioner_diag_gpu, + preconditioner_reciprocal_diag=u_dilu_reciprocal_diag_gpu, + dilu_schedule=u_dilu_level_schedule, + ) + u_first_intermediate_residual = np.asarray(u_first_intermediate_residual_gpu.to_numpy(), dtype=np.float64) + u_first_second_preconditioned = np.asarray(u_first_second_preconditioned_gpu.to_numpy(), dtype=np.float64) + u_first_operator_second_preconditioned = np.asarray(u_first_operator_second_preconditioned_gpu.to_numpy(), dtype=np.float64) + u_first_residual_after_omega = np.asarray(u_first_residual_after_omega_gpu.to_numpy(), dtype=np.float64) + u_first_solution_after_omega = np.asarray(u_first_solution_after_omega_gpu.to_numpy(), dtype=np.float64) u_gpu_dilu_vs_reference = np.abs(u_gpu_dilu_preconditioned - u_openfoam_dilu_preconditioned) u_diagonal_vs_reference = np.abs(u_diagonal_preconditioned - u_openfoam_dilu_preconditioned) u_dilu_preconditioner_diagnostic = { @@ -2122,6 +2408,8 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P min_iterations=DEFAULT_MOMENTUM_PBICGSTAB_MIN_ITERATIONS, normalization_factors=u_normalization_factors, ) + if trace_substitution_arrays: + apply_trace_substitution("solver.solve_UEqn.field_after", u_solved_gpu, np.asarray(u_solved_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_rans_momentum_hbyA_by_cell(n_cells, offdiag_cell_offsets_gpu, offdiag_cell_faces_gpu, offdiag_cell_sides_gpu, owner_gpu, neighbour_gpu, u_upper_gpu, u_lower_gpu, u_solved_gpu, u_source_gpu, HbyA_gpu) gpu_rans_momentum_hbyA_finish(n_cells, u_boundary_diag_candidate_gpu, HbyA_gpu) gpu_rans_pressure_inputs(n_cells, u_boundary_diag_candidate_gpu, cell_volumes_gpu, rAU_gpu, H1_gpu) @@ -2129,11 +2417,20 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P gpu_rans_consistent_rAtU(n_cells, rAU_gpu, H1_gpu, DEFAULT_SIMPLE_CONSISTENT_RATU_FACTOR, rAtU_gpu) gpu_rans_surface_flux_from_cells(n_internal_faces, owner_gpu, neighbour_gpu, HbyA_gpu, sf_gpu, face_weights_gpu, phiHbyA_gpu) gpu_rans_consistent_phiHbyA_correction(n_internal_faces, owner_gpu, neighbour_gpu, rAU_gpu, rAtU_gpu, p_gpu, cell_centres_gpu, sf_gpu, mag_sf_gpu, face_weights_gpu, phiHbyA_gpu) + if trace_substitution_arrays: + apply_trace_substitution("pressure_inputs.rAU", rAU_gpu, np.asarray(rAU_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("pressure_inputs.rAtU", rAtU_gpu, np.asarray(rAtU_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("pressure_inputs.HbyA", HbyA_gpu, np.asarray(HbyA_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("pressure_inputs.phiHbyA", phiHbyA_gpu, np.asarray(phiHbyA_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_rans_pressure_assembly(n_cells, p_gpu, p_diag_gpu, p_source_gpu) gpu_rans_pressure_laplacian_coefficients(n_internal_faces, owner_gpu, neighbour_gpu, rAtU_gpu, cell_centres_gpu, sf_gpu, mag_sf_gpu, face_weights_gpu, p_diag_gpu, p_upper_gpu) gpu_rans_pressure_source_from_flux(n_internal_faces, owner_gpu, neighbour_gpu, phiHbyA_gpu, p_source_gpu) if n_pressure_boundary_faces: gpu_rans_pressure_source_from_boundary_flux(n_pressure_boundary_faces, boundary_face_cells_gpu, boundary_phiHbyA_gpu, p_source_gpu) + if trace_substitution_arrays: + apply_trace_substitution("matrix_operator.pEqn.diag", p_diag_gpu, np.asarray(p_diag_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("matrix_operator.pEqn.upper", p_upper_gpu, np.asarray(p_upper_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) + apply_trace_substitution("matrix_operator.pEqn.source", p_source_gpu, np.asarray(p_source_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_rans_negate_scalar_field(n_cells, p_diag_gpu, p_solve_diag_gpu) if n_pressure_mixed_faces: gpu_rans_pressure_mixed_solver_diag( @@ -2173,6 +2470,58 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P p_dic_reciprocal_diag_np, "gpu_stages.solve_pEqn.dic_reciprocal_diag", ) + gpu_ldu_matvec_scalar_symmetric_faces(n_cells, n_internal_faces, owner_gpu, neighbour_gpu, p_solve_upper_gpu, p_solve_diag_gpu, p_gpu, p_operator_gpu) + gpu_zero_scalar_accumulator(p_rr_gpu) + gpu_zero_scalar_accumulator(p_denominator_gpu) + gpu_pcg_initialize_residual_scalar(n_cells, p_solve_source_gpu, p_operator_gpu, p_gpu, p_solved_gpu, p_residual_gpu, p_rr_gpu) + gpu_dic_apply_scalar_symmetric_levels(p_dic_level_schedule, owner_gpu, neighbour_gpu, p_solve_upper_gpu, p_dic_reciprocal_diag_gpu, p_residual_gpu, p_work_gpu) + gpu_copy_scalar(n_cells, p_work_gpu, p_direction_gpu) + gpu_cg_dot_scalar(n_cells, p_residual_gpu, p_work_gpu, p_denominator_gpu) + qd.sync() + p_initial_residual = np.asarray(p_residual_gpu.to_numpy(), dtype=np.float64) + p_initial_residual_squared = float(np.asarray(p_rr_gpu.to_numpy())[0]) + p_dic_preconditioned = np.asarray(p_work_gpu.to_numpy(), dtype=np.float64) + p_first_rho = float(np.asarray(p_denominator_gpu.to_numpy())[0]) + gpu_ldu_matvec_scalar_symmetric_faces(n_cells, n_internal_faces, owner_gpu, neighbour_gpu, p_solve_upper_gpu, p_solve_diag_gpu, p_direction_gpu, p_operator_gpu) + gpu_zero_scalar_accumulator(p_denominator_gpu) + gpu_cg_dot_scalar(n_cells, p_direction_gpu, p_operator_gpu, p_denominator_gpu) + qd.sync() + p_operator_preconditioned_direction = np.asarray(p_operator_gpu.to_numpy(), dtype=np.float64) + p_first_denominator = float(np.asarray(p_denominator_gpu.to_numpy())[0]) + p_first_alpha = p_first_rho / p_first_denominator if abs(p_first_denominator) > 1.0e-300 else None + p_linear_solver_trace = { + "schema_version": 1, + "solver": "gpu_symmetric_ldu_pcg", + "field": "p", + "preconditioner": "DIC level-scheduled GPU solver candidate", + "trace_points": [ + { + "name": "initial_residual", + "step": "r0 = source - A*x0", + "stats": array_stats(p_initial_residual), + "residual_squared": finite_float(p_initial_residual_squared), + "l2": finite_float(np.linalg.norm(p_initial_residual.reshape(-1))), + }, + { + "name": "preconditioned_residual", + "step": "z0 = M^-1 r0", + "candidate": "gpu_dic_apply_scalar_symmetric_levels", + "candidate_stats": array_stats(p_dic_preconditioned), + }, + { + "name": "matvec_preconditioned_search_direction", + "step": "A*z0 before alpha", + "stats": array_stats(p_operator_preconditioned_direction), + }, + { + "name": "first_scalar_reductions", + "step": "wArA0, wApA0, alpha0", + "rho": finite_float(p_first_rho), + "denominator": finite_float(p_first_denominator), + "alpha": finite_float(p_first_alpha), + }, + ], + } p_solve_passes: list[dict[str, Any]] = [] p_solve_initial_gpu = p_gpu p_solve_initial_np = p_np @@ -2221,7 +2570,11 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "solve_passes": p_solve_passes, "total_iterations": sum(int(pass_report["iterations"]) for pass_report in p_solve_passes), } + if trace_substitution_arrays: + apply_trace_substitution("solver.solve_pEqn.p", p_solved_gpu, np.asarray(p_solved_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_rans_pressure_flux_correction(n_internal_faces, owner_gpu, neighbour_gpu, p_solved_gpu, p_upper_gpu, phiHbyA_gpu, phi_solved_gpu) + if trace_substitution_arrays: + apply_trace_substitution("solver.solve_pEqn.phi", phi_solved_gpu, np.asarray(phi_solved_gpu.to_numpy(), dtype=np.float64), trace_substitution_arrays, applied_trace_substitutions) gpu_rans_final_correction(n_cells, HbyA_gpu, p_solved_gpu, p_gpu, u_final_gpu, p_final_gpu) gpu_rans_pressure_velocity_correction(n_internal_faces, owner_gpu, neighbour_gpu, rAtU_gpu, p_solved_gpu, cell_centres_gpu, sf_gpu, u_final_gpu) gpu_rans_turbulence_update(n_cells, nut_gpu, k_gpu, omega_gpu, nut_out_gpu, k_out_gpu, omega_out_gpu) @@ -2238,8 +2591,14 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_boundary_relax_add = np.asarray(u_boundary_relax_add_gpu.to_numpy()) u_boundary_relax_subtract = np.asarray(u_boundary_relax_subtract_gpu.to_numpy()) u_source = np.asarray(u_source_gpu.to_numpy()) + u_solve_diag = np.repeat(u_boundary_diag_candidate[:, None], 3, axis=1) u_unrelaxed_diag = np.asarray(u_unrelaxed_diag_gpu.to_numpy()) u_unrelaxed_source = np.asarray(u_unrelaxed_source_gpu.to_numpy()) + u_term_div_source = np.asarray(u_term_div_source_gpu.to_numpy()) + u_term_stress_source = np.asarray(u_term_stress_source_gpu.to_numpy()) + u_term_stress_wall_diffusion_source = np.asarray(u_term_stress_wall_diffusion_source_gpu.to_numpy()) + u_term_stress_internal_dev_tau_source = np.asarray(u_term_stress_internal_dev_tau_source_gpu.to_numpy()) + u_term_stress_wall_dev_tau_source = np.asarray(u_term_stress_wall_dev_tau_source_gpu.to_numpy()) u_matrix_source = np.asarray(u_matrix_source_gpu.to_numpy()) u_upper = np.asarray(u_upper_gpu.to_numpy()) u_lower = np.asarray(u_lower_gpu.to_numpy()) @@ -2251,6 +2610,9 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P p_diag = np.asarray(p_diag_gpu.to_numpy()) p_source = np.asarray(p_source_gpu.to_numpy()) p_upper = np.asarray(p_upper_gpu.to_numpy()) + p_solve_diag = np.asarray(p_solve_diag_gpu.to_numpy()) + p_solve_source = np.asarray(p_solve_source_gpu.to_numpy()) + p_solve_upper = np.asarray(p_solve_upper_gpu.to_numpy()) u_solved = np.asarray(u_solved_gpu.to_numpy()) u_residual = np.asarray(u_residual_gpu.to_numpy()) p_solved = np.asarray(p_solved_gpu.to_numpy()) @@ -2280,6 +2642,20 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "source": gpu_array_output("UEqn.unrelaxed_source", u_unrelaxed_source, kernel="gpu_copy_vector"), } UEqn["solve_rhs_source"] = gpu_array_output("UEqn.solve_rhs_source", u_source, kernel="gpu_rans_add_vector_field") + UEqn["solve_matrix_before"] = { + "diag": gpu_array_output("UEqn.solve_matrix_before.diag", u_diag, kernel="gpu_rans_momentum_equation_relaxation"), + "upper": gpu_array_output("UEqn.solve_matrix_before.upper", u_upper, kernel="gpu_rans_momentum_diffusion_coefficients"), + "lower": gpu_array_output("UEqn.solve_matrix_before.lower", u_lower, kernel="gpu_rans_momentum_convection_coefficients"), + "source": gpu_array_output("UEqn.solve_matrix_before.source", u_matrix_source, kernel="gpu_copy_vector"), + "semantics": "OpenFOAM-visible fvMatrix before solve(UEqn == -grad(p)); boundary-source and pressure-gradient RHS are reported separately", + } + UEqn["solver_operands"] = { + "diag": gpu_array_output("UEqn.solver_operands.diag", u_boundary_diag_candidate, kernel="gpu_rans_momentum_boundary_internal_diag"), + "upper": gpu_array_output("UEqn.solver_operands.upper", u_upper, kernel="gpu_rans_momentum_diffusion_coefficients"), + "lower": gpu_array_output("UEqn.solver_operands.lower", u_lower, kernel="gpu_rans_momentum_convection_coefficients"), + "source": gpu_array_output("UEqn.solver_operands.source", u_source, kernel="gpu_rans_add_vector_field"), + "semantics": "GPU linear-solver operands after hidden fvMatrix solve-time boundary diagonal/source and explicit pressure-gradient RHS", + } UEqn["boundary_diag_candidate"] = gpu_array_output("UEqn.boundary_diag_candidate", u_boundary_diag_candidate, kernel="gpu_rans_momentum_boundary_internal_diag") UEqn["boundary_diag_coeff"] = gpu_array_output("UEqn.boundary_diag_coeff", u_boundary_diag_coeff, kernel="gpu_rans_momentum_boundary_relaxation_coefficients") UEqn["boundary_relax"] = { @@ -2306,6 +2682,36 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "boundary_face_count": n_pressure_boundary_faces, "mixed_boundary_face_count": n_pressure_mixed_faces, } + pEqn["solve_matrix_before"] = { + "diag": gpu_array_output("pEqn.solve_matrix_before.diag", p_diag, kernel="gpu_rans_pressure_assembly"), + "upper": gpu_array_output("pEqn.solve_matrix_before.upper", p_upper, kernel="gpu_rans_pressure_laplacian_coefficients"), + "source": gpu_array_output("pEqn.solve_matrix_before.source", p_source, kernel="gpu_rans_pressure_source_from_flux"), + "semantics": "OpenFOAM-visible fvMatrix before solve(); SPD transform and hidden mixed-boundary diagonal are reported separately", + } + pEqn["solver_operands"] = { + "diag": gpu_array_output("pEqn.solver_operands.spd_diag", p_solve_diag, kernel="gpu_rans_negate_scalar_field+gpu_rans_pressure_mixed_solver_diag"), + "upper": gpu_array_output("pEqn.solver_operands.spd_upper", p_solve_upper, kernel="gpu_rans_negate_scalar_field"), + "source": gpu_array_output("pEqn.solver_operands.spd_source", p_solve_source, kernel="gpu_rans_negate_scalar_field"), + "operator_sign_convention": "negative_openfoam_laplacian_to_spd", + } + momentum_stage_kernels = [ + "gpu_rans_momentum_assembly", + "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", + "gpu_rans_momentum_wall_diffusion_coefficients", + "gpu_rans_momentum_convection_coefficients", + "gpu_rans_momentum_bounded_convection_sp_internal", + "gpu_rans_momentum_bounded_convection_sp_boundary", + "gpu_rans_momentum_convection_boundary_coefficients", + "gpu_rans_momentum_gauss_grad_u_by_cell", + "gpu_rans_momentum_internal_dev_tau_source", + "gpu_rans_momentum_wall_dev_tau_source", + "gpu_rans_momentum_linear_upwind_source", + "gpu_rans_momentum_diag_from_offdiag_by_cell", + "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", + "gpu_rans_momentum_equation_relaxation", + "gpu_rans_momentum_pressure_gradient_source", + "gpu_rans_momentum_pressure_boundary_source", + ] stages = [ gpu_stage_result( "momentum_transport_predict", @@ -2314,19 +2720,43 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P ), gpu_stage_result( "assemble_momentum_terms", - {"terms": [{"name": "gpu_momentum_ddt_diagonal", "kernel_entrypoint": "gpu_rans_momentum_assembly"}, {"name": "gpu_momentum_laminar_turbulent_diffusion", "kernel_entrypoint": "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "laminar_nu": laminar_nu}, {"name": "gpu_momentum_wall_diffusion", "kernel_entrypoint": "gpu_rans_momentum_wall_diffusion_coefficients", "laminar_nu": laminar_nu, "boundary_face_count": n_momentum_wall_faces, "patch_types": ["noSlip"]}, {"name": "gpu_momentum_bounded_upwind_convection", "kernel_entrypoint": "gpu_rans_momentum_convection_coefficients", "source": "fields.phi.internal"}, {"name": "gpu_momentum_bounded_convection_sp", "kernel_entrypoints": ["gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary"], "source": "-fvm::Sp(fvc::surfaceIntegrate(phi), U)"}, {"name": "gpu_momentum_convection_boundary_coefficients", "kernel_entrypoint": "gpu_rans_momentum_convection_boundary_coefficients", "source": "U.boundaryField valueInternalCoeffs/valueBoundaryCoeffs"}, {"name": "gpu_momentum_linear_upwind_correction", "kernel_entrypoints": ["gpu_rans_momentum_gauss_grad_u_internal", "gpu_rans_momentum_gauss_grad_u_boundary", "gpu_rans_momentum_linear_upwind_source"], "source": "bounded Gauss linearUpwind grad(U) explicit correction"}, {"name": "gpu_momentum_equation_relaxation", "kernel_entrypoints": ["gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation"], "alpha": DEFAULT_MOMENTUM_RELAXATION_ALPHA}, {"name": "gpu_momentum_pressure_gradient_source", "kernel_entrypoint": "gpu_rans_momentum_pressure_gradient_source", "source": "-fvc::grad(p)"}]}, - kernels=["gpu_rans_momentum_assembly", "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "gpu_rans_momentum_wall_diffusion_coefficients", "gpu_rans_momentum_convection_coefficients", "gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary", "gpu_rans_momentum_convection_boundary_coefficients", "gpu_rans_zero_tensor_field", "gpu_rans_momentum_gauss_grad_u_internal", "gpu_rans_momentum_gauss_grad_u_boundary", "gpu_rans_momentum_gauss_grad_u_finish", "gpu_rans_momentum_linear_upwind_source", "gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation", "gpu_rans_momentum_pressure_gradient_source", "gpu_rans_momentum_pressure_boundary_source"], + { + "terms": [ + {"name": "gpu_momentum_ddt_diagonal", "kernel_entrypoint": "gpu_rans_momentum_assembly"}, + {"name": "gpu_momentum_laminar_turbulent_diffusion", "kernel_entrypoint": "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "laminar_nu": laminar_nu}, + {"name": "gpu_momentum_wall_diffusion", "kernel_entrypoint": "gpu_rans_momentum_wall_diffusion_coefficients", "laminar_nu": laminar_nu, "boundary_face_count": n_momentum_wall_faces, "patch_types": ["noSlip"]}, + {"name": "gpu_momentum_bounded_upwind_convection", "kernel_entrypoint": "gpu_rans_momentum_convection_coefficients", "source": "fields.phi.internal"}, + {"name": "gpu_momentum_bounded_convection_sp", "kernel_entrypoints": ["gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary"], "source": "-fvm::Sp(fvc::surfaceIntegrate(phi), U)"}, + {"name": "gpu_momentum_convection_boundary_coefficients", "kernel_entrypoint": "gpu_rans_momentum_convection_boundary_coefficients", "source": "U.boundaryField valueInternalCoeffs/valueBoundaryCoeffs"}, + {"name": "gpu_momentum_gauss_grad_u", "kernel_entrypoint": "gpu_rans_momentum_gauss_grad_u_by_cell", "source": "OpenFOAM face-order Gauss grad(U)"}, + {"name": "gpu_momentum_dev_tau_correction", "kernel_entrypoints": ["gpu_rans_momentum_internal_dev_tau_source", "gpu_rans_momentum_wall_dev_tau_source"], "source": "momentumTransport->divDevSigma(U) explicit correction"}, + {"name": "gpu_momentum_linear_upwind_correction", "kernel_entrypoints": ["gpu_rans_momentum_gauss_grad_u_by_cell", "gpu_rans_momentum_linear_upwind_source"], "source": "bounded Gauss linearUpwind grad(U) explicit correction"}, + {"name": "gpu_momentum_equation_relaxation", "kernel_entrypoints": ["gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation"], "alpha": DEFAULT_MOMENTUM_RELAXATION_ALPHA}, + {"name": "gpu_momentum_pressure_gradient_source", "kernel_entrypoint": "gpu_rans_momentum_pressure_gradient_source", "source": "-fvc::grad(p)"}, + ] + }, + kernels=momentum_stage_kernels, + ), + gpu_stage_result( + "assemble_UEqn", + {"UEqn": UEqn, "relaxation": {"alpha": DEFAULT_MOMENTUM_RELAXATION_ALPHA, "kernel_entrypoints": ["gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation"]}}, + kernels=momentum_stage_kernels, ), - gpu_stage_result("assemble_UEqn", {"UEqn": UEqn, "relaxation": {"alpha": DEFAULT_MOMENTUM_RELAXATION_ALPHA, "kernel_entrypoints": ["gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation"]}}, kernels=["gpu_rans_momentum_assembly", "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "gpu_rans_momentum_wall_diffusion_coefficients", "gpu_rans_momentum_convection_coefficients", "gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary", "gpu_rans_momentum_convection_boundary_coefficients", "gpu_rans_zero_tensor_field", "gpu_rans_momentum_gauss_grad_u_internal", "gpu_rans_momentum_gauss_grad_u_boundary", "gpu_rans_momentum_gauss_grad_u_finish", "gpu_rans_momentum_linear_upwind_source", "gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation", "gpu_rans_momentum_pressure_gradient_source", "gpu_rans_momentum_pressure_boundary_source"]), gpu_stage_result( "solve_UEqn", { "performance": {"solver_name": "gpu_asymmetric_ldu_pbicgstab", "field_name": "U", **u_solve_performance}, "rhs_source": gpu_array_output("UEqn.solve_rhs_source", u_source, kernel="gpu_rans_add_vector_field"), + "solve_diag": gpu_array_output("UEqn.solve_diag", u_solve_diag, kernel="gpu_rans_momentum_boundary_internal_diag"), + "solve_source": gpu_array_output("UEqn.solve_source", u_source, kernel="gpu_rans_add_vector_field"), + "level_preconditioned_residual": gpu_array_output("UEqn.level_preconditioned_residual", u_gpu_level_preconditioned, kernel="gpu_dilu_apply_vector_asymmetric_levels"), + "matrix_before": UEqn["solve_matrix_before"], + "solver_operands": UEqn["solver_operands"], "field_after": gpu_array_output("U", u_solved, kernel="gpu_bicgstab_update_solution_residual_vector_preconditioned"), "residual": gpu_array_output("U_residual", u_residual, kernel="gpu_bicgstab_update_solution_residual_vector_preconditioned"), "preconditioner_diagnostic": u_dilu_preconditioner_diagnostic, "linear_solver_trace": u_linear_solver_trace, + "trace_substitutions": {"requested": sorted(trace_substitution_arrays), "applied": applied_trace_substitutions}, }, kernels=["gpu_extract_vector_component", "gpu_scatter_vector_component", "gpu_ldu_matvec_vector_asymmetric_diag", "gpu_ldu_matvec_vector_asymmetric_face_accumulate", "gpu_bicgstab_initialize_vector", "gpu_bicgstab_dot_vector", "gpu_bicgstab_update_direction_vector", "gpu_bicgstab_precondition_vector", "gpu_dilu_apply_vector_asymmetric_faces", "gpu_dilu_forward_level_vector", "gpu_dilu_backward_level_vector", "gpu_bicgstab_update_intermediate_vector_preconditioned", "gpu_bicgstab_update_solution_residual_vector_preconditioned"], changed_fields=["U"], @@ -2372,7 +2802,11 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "performance": {"solver_name": "gpu_symmetric_ldu_pcg", "field_name": "p", "operator_transform": "negative_openfoam_laplacian_to_spd", "non_orthogonal_loop": "OpenFOAM SIMPLE nNonOrthogonalCorrectors=3", **p_solve_performance}, "p": gpu_array_output("p", p_solved, kernel="gpu_pcg_update_solution_residual_only_scalar"), "residual": gpu_array_output("p_residual", p_residual, kernel="gpu_pcg_update_solution_residual_only_scalar"), + "matrix_before": pEqn["solve_matrix_before"], + "solver_operands": pEqn["solver_operands"], "phi": gpu_array_output("phi", phi_solved, kernel="gpu_rans_pressure_flux_correction"), + "linear_solver_trace": p_linear_solver_trace, + "trace_substitutions": {"requested": sorted(trace_substitution_arrays), "applied": applied_trace_substitutions}, }, kernels=["gpu_rans_negate_scalar_field", "gpu_rans_pressure_mixed_solver_diag", "gpu_ldu_matvec_scalar_symmetric_diag", "gpu_ldu_matvec_scalar_symmetric_face_accumulate", "gpu_pcg_initialize_residual_scalar", "gpu_dic_forward_level_scalar", "gpu_dic_backward_level_scalar", "gpu_cg_dot_scalar", "gpu_pcg_update_solution_residual_only_scalar", "gpu_pcg_update_direction_scalar"], changed_fields=["p", "phi"], @@ -2445,10 +2879,21 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P diagnostic_artifacts = write_numeric_artifact_file( diagnostic_artifact_path, { + "matrix_terms.UEqn.ddt.source": np.zeros_like(u_np), + "matrix_terms.UEqn.div.source": u_term_div_source, + "matrix_terms.UEqn.divDevSigma.source": u_term_stress_source, + "matrix_terms.UEqn.divDevSigma.wall_dev_tau.source": u_term_stress_wall_dev_tau_source, + "candidate_diagnostics.UEqn.divDevSigma.wall_diffusion.source": u_term_stress_wall_diffusion_source, + "candidate_diagnostics.UEqn.divDevSigma.internal_dev_tau.source": u_term_stress_internal_dev_tau_source, + "candidate_diagnostics.UEqn.divDevSigma.wall_dev_tau.source": u_term_stress_wall_dev_tau_source, + "matrix_terms.UEqn.fvModels.source": np.zeros_like(u_np), + "matrix_terms.UEqn.MRF_DDt.field": np.zeros_like(u_np), "matrix_operator.UEqn.diag": u_diag, "matrix_operator.UEqn.upper": u_upper, "matrix_operator.UEqn.lower": u_lower, "matrix_operator.UEqn.source": u_matrix_source, + "matrix_operator.UEqn.unrelaxed_diag": u_unrelaxed_diag, + "matrix_operator.UEqn.unrelaxed_source": u_unrelaxed_source, "matrix_operator.UEqn.psi": u_np, "matrix_operator.pEqn.diag": p_diag, "matrix_operator.pEqn.upper": p_upper, @@ -2458,8 +2903,39 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "pressure_inputs.phiHbyA": phiHbyA, "pressure_inputs.rAU": rAU, "pressure_inputs.rAtU": rAtU, + "solver.solve_UEqn.matrix_before.diag": u_diag, + "solver.solve_UEqn.matrix_before.upper": u_upper, + "solver.solve_UEqn.matrix_before.lower": u_lower, + "solver.solve_UEqn.matrix_before.source": u_matrix_source, + "solver.solve_UEqn.solve_diag": u_solve_diag, + "solver.solve_UEqn.solve_source": u_source, + "solver.solve_UEqn.initial_residual": u_initial_residual, + "solver.solve_UEqn.diagonal_preconditioned_residual": u_diagonal_preconditioned, + "solver.solve_UEqn.preconditioned_residual": u_gpu_dilu_preconditioned, + "solver.solve_UEqn.level_preconditioned_residual": u_gpu_level_preconditioned, + "solver.solve_UEqn.operator_preconditioned_direction": u_operator_preconditioned_direction, + "solver.solve_UEqn.first_iteration.rho": np.asarray(u_first_iteration_diagnostics["rho"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.denominator": np.asarray(u_first_iteration_diagnostics["denominator"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.alpha": np.asarray(u_first_iteration_diagnostics["alpha"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.intermediate_residual": u_first_intermediate_residual, + "solver.solve_UEqn.first_iteration.second_preconditioned_residual": u_first_second_preconditioned, + "solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual": u_first_operator_second_preconditioned, + "solver.solve_UEqn.first_iteration.omega_numerator": np.asarray(u_first_iteration_diagnostics["omega_numerator"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.omega_denominator": np.asarray(u_first_iteration_diagnostics["omega_denominator"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.omega": np.asarray(u_first_iteration_diagnostics["omega"], dtype=np.float64), + "solver.solve_UEqn.first_iteration.residual_after_omega": u_first_residual_after_omega, + "solver.solve_UEqn.first_iteration.solution_after_omega": u_first_solution_after_omega, "solver.solve_UEqn.field_after": u_solved, "solver.solve_UEqn.residual": u_residual, + "solver.solve_pEqn.matrix_before.diag": p_diag, + "solver.solve_pEqn.matrix_before.upper": p_upper, + "solver.solve_pEqn.matrix_before.source": p_source, + "solver.solve_pEqn.spd_diag": p_solve_diag, + "solver.solve_pEqn.spd_upper": p_solve_upper, + "solver.solve_pEqn.spd_source": p_solve_source, + "solver.solve_pEqn.initial_residual": p_initial_residual, + "solver.solve_pEqn.preconditioned_residual": p_dic_preconditioned, + "solver.solve_pEqn.operator_preconditioned_direction": p_operator_preconditioned_direction, "solver.solve_pEqn.p": p_solved, "solver.solve_pEqn.phi": phi_solved, "solver.solve_pEqn.residual": p_residual, @@ -2479,6 +2955,69 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P role="split", ) + + input_transfers = [ + u_transfer, + p_transfer, + phi_transfer, + nut_transfer, + k_transfer, + omega_transfer, + owner_transfer, + neighbour_transfer, + offdiag_cell_offsets_transfer, + offdiag_cell_faces_transfer, + offdiag_cell_sides_transfer, + grad_cell_offsets_transfer, + grad_cell_faces_transfer, + grad_cell_sides_transfer, + sf_transfer, + face_centres_transfer, + cell_centres_transfer, + cell_volumes_transfer, + mag_sf_transfer, + face_weights_transfer, + momentum_diffusion_coeff_transfer, + momentum_internal_diag_transfer, + momentum_internal_h1_transfer, + boundary_face_cells_transfer, + boundary_phi_transfer, + boundary_phiHbyA_transfer, + momentum_pressure_boundary_face_cells_transfer, + momentum_pressure_boundary_values_transfer, + momentum_pressure_boundary_sf_transfer, + momentum_pressure_source_transfer, + momentum_u_boundary_face_cells_transfer, + momentum_u_boundary_values_transfer, + momentum_u_boundary_sf_transfer, + momentum_u_boundary_cell_offsets_transfer, + momentum_u_boundary_cell_faces_transfer, + momentum_u_convection_face_cells_transfer, + momentum_u_convection_phi_transfer, + momentum_u_convection_internal_coeff_transfer, + momentum_u_convection_boundary_coeff_transfer, + pressure_mixed_face_cells_transfer, + pressure_mixed_u_values_transfer, + pressure_mixed_face_centres_transfer, + pressure_mixed_area_vectors_transfer, + pressure_mixed_area_magnitudes_transfer, + momentum_wall_face_cells_transfer, + momentum_wall_values_transfer, + momentum_wall_nut_transfer, + momentum_wall_face_centres_transfer, + momentum_wall_area_vectors_transfer, + momentum_wall_area_magnitudes_transfer, + momentum_wall_delta_coeffs_transfer, + momentum_wall_normals_transfer, + momentum_wall_gamma_area_vectors_transfer, + momentum_wall_trace_gamma_area_vectors_transfer, + omega_wall_cells_transfer, + omega_wall_distances_transfer, + u_preconditioner_diag_transfer, + u_dilu_reciprocal_diag_transfer, + p_preconditioner_diag_transfer, + p_dic_reciprocal_diag_transfer, + ] return { "status": "executed", "case": case, @@ -2501,8 +3040,9 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P "k": gpu_array_output("k", k_out, kernel="gpu_rans_turbulence_update"), "omega": gpu_array_output("omega", omega_out, kernel="gpu_rans_omega_wall_update"), }, + "field_objects": field_objects, "transfers": { - "inputs": [u_transfer, p_transfer, phi_transfer, nut_transfer, k_transfer, omega_transfer, owner_transfer, neighbour_transfer, offdiag_cell_offsets_transfer, offdiag_cell_faces_transfer, offdiag_cell_sides_transfer, sf_transfer, face_centres_transfer, cell_centres_transfer, cell_volumes_transfer, mag_sf_transfer, face_weights_transfer, momentum_diffusion_coeff_transfer, momentum_internal_diag_transfer, momentum_internal_h1_transfer, boundary_face_cells_transfer, boundary_phi_transfer, boundary_phiHbyA_transfer, momentum_pressure_boundary_face_cells_transfer, momentum_pressure_boundary_values_transfer, momentum_pressure_boundary_sf_transfer, momentum_pressure_source_transfer, momentum_u_boundary_face_cells_transfer, momentum_u_boundary_values_transfer, momentum_u_boundary_sf_transfer, momentum_u_convection_face_cells_transfer, momentum_u_convection_phi_transfer, momentum_u_convection_internal_coeff_transfer, momentum_u_convection_boundary_coeff_transfer, pressure_mixed_face_cells_transfer, pressure_mixed_u_values_transfer, pressure_mixed_face_centres_transfer, pressure_mixed_area_vectors_transfer, pressure_mixed_area_magnitudes_transfer, momentum_wall_face_cells_transfer, momentum_wall_values_transfer, momentum_wall_nut_transfer, momentum_wall_face_centres_transfer, momentum_wall_area_vectors_transfer, momentum_wall_area_magnitudes_transfer, omega_wall_cells_transfer, omega_wall_distances_transfer, u_preconditioner_diag_transfer, u_dilu_reciprocal_diag_transfer, p_preconditioner_diag_transfer, p_dic_reciprocal_diag_transfer], + "inputs": input_transfers, "allocations": [ u_diag_alloc, u_source_alloc, @@ -2516,6 +3056,15 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_upper_alloc, u_lower_alloc, grad_u_alloc, + u_term_div_diag_alloc, + u_term_div_source_alloc, + u_term_stress_boundary_relax_add_alloc, + u_term_stress_boundary_relax_subtract_alloc, + u_term_stress_boundary_diag_alloc, + u_term_stress_source_alloc, + u_term_stress_wall_diffusion_source_alloc, + u_term_stress_internal_dev_tau_source_alloc, + u_term_stress_wall_dev_tau_source_alloc, rAU_alloc, H1_alloc, rAtU_alloc, @@ -2530,6 +3079,12 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P u_operator_direction_alloc, u_intermediate_alloc, u_operator_intermediate_alloc, + u_level_preconditioned_alloc, + u_first_intermediate_residual_alloc, + u_first_second_preconditioned_alloc, + u_first_operator_second_preconditioned_alloc, + u_first_residual_after_omega_alloc, + u_first_solution_after_omega_alloc, u_rr_alloc, u_rho_alloc, u_denominator_alloc, @@ -2554,6 +3109,11 @@ def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: P omega_alloc, ], }, + "trace_substitutions": { + "requested": sorted(trace_substitution_arrays), + "applied": applied_trace_substitutions, + "used_cpu_fallback": False, + }, "profiler": profiler_evidence, "timing": timing, "diagnostic_artifacts": diagnostic_artifacts, @@ -2660,6 +3220,8 @@ def select_gpu_solver_backend(gpu_present: bool) -> dict[str, Any]: "failure_diagnostics:missing_kernels", "failure_diagnostics:unsupported_stage", "failure_diagnostics:numerical_mismatch", + "differential_trace:checkpoint_artifacts", + "differential_trace:reference_substitution_replay", *stage_capabilities, ], "verifier_integration": { @@ -2830,8 +3392,8 @@ def run_backend_iteration(stepper: Any, backend: Mapping[str, Any], case: Path) } -def run_gpu_split_iteration(foam: Any, stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None) -> dict[str, Any]: - gpu_run = run_gpu_solver_stage_smoke(stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path) +def run_gpu_split_iteration(foam: Any, stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None, trace_substitutions: Mapping[str, Any] | None = None) -> dict[str, Any]: + gpu_run = run_gpu_solver_stage_smoke(stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path, trace_substitutions=trace_substitutions) fields = gpu_run["field_objects"] stages = list(gpu_run["stages"]) observability = stage_observability_report( diff --git a/python/src/foam_stepper/gpu/constants.py b/python/src/foam_stepper/gpu/constants.py index 32176a4..2e96415 100644 --- a/python/src/foam_stepper/gpu/constants.py +++ b/python/src/foam_stepper/gpu/constants.py @@ -97,7 +97,7 @@ STAGE_OBSERVABILITY_GROUPS = ( ) GPU_STAGE_KERNEL_ENTRYPOINTS = { - "momentum_assembly": ["gpu_rans_momentum_assembly", "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "gpu_rans_momentum_wall_diffusion_coefficients", "gpu_rans_momentum_convection_coefficients", "gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary", "gpu_rans_momentum_convection_boundary_coefficients", "gpu_rans_momentum_convection_boundary_source", "gpu_rans_momentum_boundary_internal_diag", "gpu_rans_momentum_boundary_relaxation_coefficients", "gpu_rans_add_scalar_field", "gpu_rans_add_vector_field", "gpu_rans_zero_tensor_field", "gpu_rans_momentum_gauss_grad_u_internal", "gpu_rans_momentum_gauss_grad_u_boundary", "gpu_rans_momentum_gauss_grad_u_finish", "gpu_rans_momentum_internal_dev_tau_source", "gpu_rans_momentum_wall_dev_tau_source", "gpu_rans_momentum_linear_upwind_source", "gpu_rans_zero_scalar_field", "gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation", "gpu_rans_momentum_pressure_gradient_source", "gpu_rans_momentum_pressure_boundary_source"], + "momentum_assembly": ["gpu_rans_momentum_assembly", "gpu_rans_momentum_diffusion_coefficients_from_face_coeff", "gpu_rans_momentum_wall_diffusion_coefficients", "gpu_rans_momentum_convection_coefficients", "gpu_rans_momentum_bounded_convection_sp_internal", "gpu_rans_momentum_bounded_convection_sp_boundary", "gpu_rans_momentum_convection_boundary_coefficients", "gpu_rans_momentum_convection_boundary_source", "gpu_rans_momentum_boundary_internal_diag", "gpu_rans_momentum_boundary_relaxation_coefficients", "gpu_rans_add_scalar_field", "gpu_rans_add_vector_field", "gpu_rans_momentum_gauss_grad_u_by_cell", "gpu_rans_momentum_internal_dev_tau_source", "gpu_rans_momentum_wall_dev_tau_source", "gpu_rans_momentum_linear_upwind_source", "gpu_rans_zero_scalar_field", "gpu_rans_momentum_diag_from_offdiag_by_cell", "gpu_rans_momentum_offdiag_abs_accumulate_by_cell", "gpu_rans_momentum_equation_relaxation", "gpu_rans_momentum_pressure_gradient_source", "gpu_rans_momentum_pressure_boundary_source"], "pressure_assembly": ["gpu_rans_pressure_inputs", "gpu_copy_scalar", "gpu_rans_momentum_h1_finish", "gpu_rans_consistent_rAtU", "gpu_rans_momentum_hbyA_by_cell", "gpu_rans_momentum_hbyA_finish", "gpu_rans_surface_flux_from_cells", "gpu_rans_consistent_phiHbyA_correction", "gpu_rans_pressure_assembly", "gpu_rans_pressure_laplacian_coefficients", "gpu_rans_pressure_source_from_flux", "gpu_rans_pressure_source_from_boundary_flux"], "linear_solve_results": ["gpu_extract_vector_component", "gpu_scatter_vector_component", "gpu_ldu_matvec_vector_asymmetric_diag", "gpu_ldu_matvec_vector_asymmetric_face_accumulate", "gpu_bicgstab_initialize_vector", "gpu_bicgstab_dot_vector", "gpu_bicgstab_update_direction_vector", "gpu_bicgstab_precondition_vector", "gpu_dilu_apply_vector_asymmetric_faces", "gpu_dilu_forward_level_vector", "gpu_dilu_backward_level_vector", "gpu_bicgstab_update_intermediate_vector_preconditioned", "gpu_bicgstab_update_solution_residual_vector_preconditioned", "gpu_vector_residual_squared", "gpu_rans_negate_scalar_field", "gpu_rans_pressure_mixed_solver_diag", "gpu_ldu_matvec_scalar_symmetric_diag", "gpu_ldu_matvec_scalar_symmetric_face_accumulate", "gpu_pcg_initialize_scalar", "gpu_pcg_initialize_residual_scalar", "gpu_dic_forward_level_scalar", "gpu_dic_backward_level_scalar", "gpu_cg_dot_scalar", "gpu_pcg_update_solution_residual_scalar", "gpu_pcg_update_solution_residual_only_scalar", "gpu_pcg_update_direction_scalar"], "final_correction": ["gpu_rans_pressure_flux_correction", "gpu_rans_final_correction", "gpu_rans_pressure_velocity_correction"], diff --git a/python/src/foam_stepper/gpu/kernels.py b/python/src/foam_stepper/gpu/kernels.py index 296c22a..424533a 100644 --- a/python/src/foam_stepper/gpu/kernels.py +++ b/python/src/foam_stepper/gpu/kernels.py @@ -340,6 +340,48 @@ def gpu_rans_momentum_gauss_grad_u_finish( grad_u[cell, component, direction] *= inv_volume +@qd.kernel +def gpu_rans_momentum_gauss_grad_u_by_cell( + n_cells: int, + internal_cell_offsets: qd.types.NDArray[qd.i32, 1], + internal_cell_faces: qd.types.NDArray[qd.i32, 1], + internal_cell_sides: qd.types.NDArray[qd.i32, 1], + owner: qd.types.NDArray[qd.i32, 1], + neighbour: qd.types.NDArray[qd.i32, 1], + u_internal: qd.types.NDArray[qd.f64, 2], + sf_internal: qd.types.NDArray[qd.f64, 2], + face_weights: qd.types.NDArray[qd.f64, 1], + boundary_cell_offsets: qd.types.NDArray[qd.i32, 1], + boundary_cell_faces: qd.types.NDArray[qd.i32, 1], + u_boundary: qd.types.NDArray[qd.f64, 2], + sf_boundary: qd.types.NDArray[qd.f64, 2], + cell_volumes: qd.types.NDArray[qd.f64, 1], + grad_u: qd.types.NDArray[qd.f64, 3], +) -> None: + for cell in range(n_cells): + inv_volume = 1.0 / cell_volumes[cell] + for component in range(3): + for direction in range(3): + total: qd.f64 = 0.0 + for entry in range(internal_cell_offsets[cell], internal_cell_offsets[cell + 1]): + face = internal_cell_faces[entry] + owner_cell = owner[face] + neighbour_cell = neighbour[face] + owner_weight = face_weights[face] + face_value = ( + owner_weight * u_internal[owner_cell, component] + + (1.0 - owner_weight) * u_internal[neighbour_cell, component] + ) + signed_flux = face_value * sf_internal[face, direction] + if internal_cell_sides[entry] != 0: + signed_flux = -signed_flux + total += signed_flux + for entry in range(boundary_cell_offsets[cell], boundary_cell_offsets[cell + 1]): + face = boundary_cell_faces[entry] + total += u_boundary[face, component] * sf_boundary[face, direction] + grad_u[cell, component, direction] = total * inv_volume + + @qd.kernel def gpu_rans_momentum_internal_dev_tau_source( n_internal_faces: int, @@ -402,32 +444,21 @@ def gpu_rans_momentum_wall_dev_tau_source( n_boundary_faces: int, face_cells: qd.types.NDArray[qd.i32, 1], boundary_values: qd.types.NDArray[qd.f64, 2], - nut_boundary: qd.types.NDArray[qd.f64, 1], - laminar_nu: qd.f64, + gamma_area_vectors: qd.types.NDArray[qd.f64, 2], + trace_gamma_area_vectors: qd.types.NDArray[qd.f64, 2], + delta_coeffs: qd.types.NDArray[qd.f64, 1], + normals: qd.types.NDArray[qd.f64, 2], u_internal: qd.types.NDArray[qd.f64, 2], - cell_centres: qd.types.NDArray[qd.f64, 2], - face_centres: qd.types.NDArray[qd.f64, 2], - face_area_vectors: qd.types.NDArray[qd.f64, 2], - face_area_magnitudes: qd.types.NDArray[qd.f64, 1], grad_u: qd.types.NDArray[qd.f64, 3], source: qd.types.NDArray[qd.f64, 2], ) -> None: for face in range(n_boundary_faces): cell = face_cells[face] - mag_sf = face_area_magnitudes[face] - nx = face_area_vectors[face, 0] / mag_sf - ny = face_area_vectors[face, 1] / mag_sf - nz = face_area_vectors[face, 2] / mag_sf - dx = face_centres[face, 0] - cell_centres[cell, 0] - dy = face_centres[face, 1] - cell_centres[cell, 1] - dz = face_centres[face, 2] - cell_centres[cell, 2] - projected_delta = dx * nx + dy * ny + dz * nz - if projected_delta < 0.0: - projected_delta = -projected_delta - if projected_delta < 1.0e-300: - projected_delta = 1.0e-300 - delta_coeff = 1.0 / projected_delta - gamma = laminar_nu + nut_boundary[face] + nx = normals[face, 0] + ny = normals[face, 1] + nz = normals[face, 2] + delta_coeff = delta_coeffs[face] + corrected_grad00 = grad_u[cell, 0, 0] corrected_grad01 = grad_u[cell, 0, 1] corrected_grad02 = grad_u[cell, 0, 2] @@ -460,18 +491,27 @@ def gpu_rans_momentum_wall_dev_tau_source( corrected_grad22 += correction * nz trace = corrected_grad00 + corrected_grad11 + corrected_grad22 - dev00 = corrected_grad00 - (2.0 / 3.0) * trace - dev01 = corrected_grad01 - dev02 = corrected_grad02 - dev10 = corrected_grad10 - dev11 = corrected_grad11 - (2.0 / 3.0) * trace - dev12 = corrected_grad12 - dev20 = corrected_grad20 - dev21 = corrected_grad21 - dev22 = corrected_grad22 - (2.0 / 3.0) * trace - qd.atomic_add(source[cell, 0], gamma * (face_area_vectors[face, 0] * dev00 + face_area_vectors[face, 1] * dev10 + face_area_vectors[face, 2] * dev20)) - qd.atomic_add(source[cell, 1], gamma * (face_area_vectors[face, 0] * dev01 + face_area_vectors[face, 1] * dev11 + face_area_vectors[face, 2] * dev21)) - qd.atomic_add(source[cell, 2], gamma * (face_area_vectors[face, 0] * dev02 + face_area_vectors[face, 1] * dev12 + face_area_vectors[face, 2] * dev22)) + qd.atomic_add( + source[cell, 0], + gamma_area_vectors[face, 0] * corrected_grad00 + + gamma_area_vectors[face, 1] * corrected_grad10 + + gamma_area_vectors[face, 2] * corrected_grad20 + - trace_gamma_area_vectors[face, 0] * trace, + ) + qd.atomic_add( + source[cell, 1], + gamma_area_vectors[face, 0] * corrected_grad01 + + gamma_area_vectors[face, 1] * corrected_grad11 + + gamma_area_vectors[face, 2] * corrected_grad21 + - trace_gamma_area_vectors[face, 1] * trace, + ) + qd.atomic_add( + source[cell, 2], + gamma_area_vectors[face, 0] * corrected_grad02 + + gamma_area_vectors[face, 1] * corrected_grad12 + + gamma_area_vectors[face, 2] * corrected_grad22 + - trace_gamma_area_vectors[face, 2] * trace, + ) @@ -506,7 +546,6 @@ def gpu_rans_momentum_linear_upwind_source( qd.atomic_add(source[owner_cell, component], -flux_correction) qd.atomic_add(source[neighbour_cell, component], flux_correction) - @qd.kernel def gpu_rans_zero_scalar_field( n_cells: int, @@ -1103,6 +1142,7 @@ __all__ = [ "gpu_rans_momentum_gauss_grad_u_internal", "gpu_rans_momentum_gauss_grad_u_boundary", "gpu_rans_momentum_gauss_grad_u_finish", + "gpu_rans_momentum_gauss_grad_u_by_cell", "gpu_rans_momentum_internal_dev_tau_source", "gpu_rans_momentum_wall_dev_tau_source", "gpu_rans_momentum_linear_upwind_source", diff --git a/python/src/foam_stepper/gpu/linear_solve.py b/python/src/foam_stepper/gpu/linear_solve.py index c65c3e4..3d59d45 100644 --- a/python/src/foam_stepper/gpu/linear_solve.py +++ b/python/src/foam_stepper/gpu/linear_solve.py @@ -452,6 +452,8 @@ def gpu_ldu_pcg_scalar_symmetric_faces( residual_l1_accumulator = qd.ndarray(qd.f64, shape=(1,)) if normalization_factor is not None and normalization_factor > 0.0 else None normalized_residual = None initial_normalized_residual = None + trace_limit = 5 + iteration_trace: list[dict[str, float | int | None]] = [] if residual_l1_accumulator is not None: gpu_zero_scalar_accumulator(residual_l1_accumulator) gpu_scalar_abs_sum(n_cells, residual, residual_l1_accumulator) @@ -475,6 +477,14 @@ def gpu_ldu_pcg_scalar_symmetric_faces( if not np.isfinite(denominator_value) or abs(denominator_value) <= 1.0e-300: break alpha = rho_value / denominator_value + trace_entry: dict[str, float | int | None] = { + "iteration": performed_iterations, + "rho": rho_value, + "denominator": denominator_value, + "alpha": alpha, + "residual_squared_before": rr_value, + "normalized_residual_before": normalized_residual, + } gpu_zero_scalar_accumulator(residual_squared) gpu_zero_scalar_accumulator(denominator) if use_dic_levels: @@ -492,6 +502,9 @@ def gpu_ldu_pcg_scalar_symmetric_faces( gpu_scalar_abs_sum(n_cells, residual, residual_l1_accumulator) qd.sync() next_normalized_residual = float(np.asarray(residual_l1_accumulator.to_numpy())[0]) / float(normalization_factor) + trace_entry["residual_squared_after"] = next_rr_value + trace_entry["rho_after"] = next_rho_value + trace_entry["normalized_residual_after"] = next_normalized_residual performed_iterations += 1 if residual_l1_accumulator is not None: if performed_iterations >= min_iterations and next_normalized_residual is not None and next_normalized_residual <= float(residual_tolerance or 0.0): @@ -504,10 +517,15 @@ def gpu_ldu_pcg_scalar_symmetric_faces( rho_value = next_rho_value break beta = next_rho_value / rho_value if rho_value != 0.0 else 0.0 + trace_entry["beta"] = beta + if len(iteration_trace) < trace_limit: + iteration_trace.append(trace_entry) gpu_pcg_update_direction_scalar(n_cells, beta, preconditioned_residual, direction) rr_value = next_rr_value rho_value = next_rho_value normalized_residual = next_normalized_residual + if performed_iterations and len(iteration_trace) < trace_limit and "trace_entry" in locals() and trace_entry not in iteration_trace: + iteration_trace.append(trace_entry) preconditioner_name = "diagonal_jacobi" if preconditioner_diag is None else "dic_reciprocal_diagonal" if use_dic_levels: @@ -528,6 +546,8 @@ def gpu_ldu_pcg_scalar_symmetric_faces( "normalization_factor": normalization_factor, "min_iterations": min_iterations, "preconditioner": preconditioner_name, + "iteration_trace": iteration_trace, + "trace_limit": trace_limit, } if use_dic_levels and dic_schedule is not None: report["preconditioner_schedule"] = dict(dic_schedule.metadata) @@ -661,6 +681,7 @@ def gpu_ldu_matvec_vector_asymmetric_face_accumulate( qd.atomic_add(out[neighbour_cell, 2], lower_coeff * current[owner_cell, 2]) + def gpu_ldu_matvec_vector_asymmetric_faces( n_cells: int, n_internal_faces: int, @@ -1484,6 +1505,8 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( residual_l1_accumulator = qd.ndarray(qd.f64, shape=(1,)) if normalization_factor is not None and normalization_factor > 0.0 else None normalized_residual = None initial_normalized_residual = None + trace_limit = 5 + iteration_trace: list[dict[str, float | int | None]] = [] if residual_l1_accumulator is not None: gpu_zero_scalar_accumulator(residual_l1_accumulator) gpu_vector_abs_sum(n_cells, residual, residual_l1_accumulator) @@ -1511,6 +1534,13 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( break else: beta = (rho_new / rho_old) * (alpha / omega) + trace_entry: dict[str, float | int | None] = { + "iteration": performed_iterations, + "rho": rho_new, + "beta": beta, + "residual_squared_before": residual_value, + "normalized_residual_before": normalized_residual, + } gpu_bicgstab_update_direction_vector(n_cells, beta, omega, residual, direction, operator_direction) if use_dilu_levels: @@ -1545,6 +1575,8 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( if not np.isfinite(denominator_value) or abs(denominator_value) <= 1.0e-300: break alpha = rho_new / denominator_value + trace_entry["denominator"] = denominator_value + trace_entry["alpha"] = alpha gpu_zero_scalar_accumulator(residual_squared) gpu_bicgstab_update_intermediate_vector_preconditioned( @@ -1565,6 +1597,8 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( gpu_vector_abs_sum(n_cells, intermediate, residual_l1_accumulator) qd.sync() intermediate_normalized_residual = float(np.asarray(residual_l1_accumulator.to_numpy())[0]) / float(normalization_factor) + trace_entry["intermediate_residual_squared"] = intermediate_residual + trace_entry["intermediate_normalized_residual"] = intermediate_normalized_residual performed_iterations += 1 if residual_l1_accumulator is not None: if performed_iterations >= min_iterations and intermediate_normalized_residual is not None and intermediate_normalized_residual <= float(residual_tolerance or 0.0): @@ -1599,6 +1633,9 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( if not np.isfinite(omega_numerator_value) or not np.isfinite(omega_denominator_value) or abs(omega_denominator_value) <= 1.0e-300: break omega = omega_numerator_value / omega_denominator_value + trace_entry["omega_numerator"] = omega_numerator_value + trace_entry["omega_denominator"] = omega_denominator_value + trace_entry["omega"] = omega gpu_zero_scalar_accumulator(residual_squared) gpu_bicgstab_update_solution_residual_vector_preconditioned( @@ -1618,7 +1655,13 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( gpu_vector_abs_sum(n_cells, residual, residual_l1_accumulator) qd.sync() normalized_residual = float(np.asarray(residual_l1_accumulator.to_numpy())[0]) / float(normalization_factor) + trace_entry["residual_squared_after"] = residual_value + trace_entry["normalized_residual_after"] = normalized_residual + if len(iteration_trace) < trace_limit: + iteration_trace.append(trace_entry) rho_old = rho_new + if performed_iterations and len(iteration_trace) < trace_limit and "trace_entry" in locals() and trace_entry not in iteration_trace: + iteration_trace.append(trace_entry) preconditioner_name = "diagonal_jacobi" if preconditioner_diag is None else "dilu_reciprocal_diagonal" if use_dilu_levels: @@ -1638,12 +1681,218 @@ def gpu_ldu_pbicgstab_vector_asymmetric_faces( "normalization_factor": normalization_factor, "min_iterations": min_iterations, "preconditioner": preconditioner_name, + "iteration_trace": iteration_trace, + "trace_limit": trace_limit, } if use_dilu_levels and dilu_schedule is not None: report["preconditioner_schedule"] = dict(dilu_schedule.metadata) return report +def gpu_ldu_pbicgstab_vector_asymmetric_first_iteration_components( + n_cells: int, + n_internal_faces: int, + owner: Any, + neighbour: Any, + upper: Any, + lower: Any, + diag: Any, + source: Any, + initial: Any, + intermediate_residual_out: Any, + second_preconditioned_residual_out: Any, + operator_second_preconditioned_residual_out: Any, + residual_after_omega_out: Any, + solution_after_omega_out: Any, + residual_squared: Any, + rho: Any, + denominator: Any, + omega_numerator: Any, + omega_denominator: Any, + preconditioner_diag: Any | None = None, + preconditioner_reciprocal_diag: Any | None = None, + dilu_schedule: GpuLduLevelSchedule | None = None, +) -> dict[str, Any]: + """Materialize OpenFOAM-style first-iteration PBiCGStab diagnostics.""" + + component_source = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_initial = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_solution = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_residual = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_shadow = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_direction = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_operator_direction = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_preconditioned_direction = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_intermediate_residual = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_second_preconditioned_residual = qd.ndarray(qd.f64, shape=(n_cells, 3)) + component_operator_second_preconditioned_residual = qd.ndarray(qd.f64, shape=(n_cells, 3)) + rho_values = np.zeros(3, dtype=np.float64) + denominator_values = np.zeros(3, dtype=np.float64) + alpha_values = np.zeros(3, dtype=np.float64) + omega_numerator_values = np.zeros(3, dtype=np.float64) + omega_denominator_values = np.zeros(3, dtype=np.float64) + omega_values = np.zeros(3, dtype=np.float64) + intermediate_residual_squared_values = np.zeros(3, dtype=np.float64) + residual_after_omega_squared_values = np.zeros(3, dtype=np.float64) + precond_diag = diag if preconditioner_diag is None else preconditioner_diag + use_dilu_levels = dilu_schedule is not None and preconditioner_reciprocal_diag is not None + + for component in range(3): + gpu_extract_vector_component(n_cells, component, source, component_source) + gpu_extract_vector_component(n_cells, component, initial, component_initial) + gpu_ldu_matvec_vector_asymmetric_faces( + n_cells, + n_internal_faces, + owner, + neighbour, + upper, + lower, + diag, + component_initial, + component_operator_second_preconditioned_residual, + ) + gpu_zero_scalar_accumulator(residual_squared) + gpu_bicgstab_initialize_vector( + n_cells, + component_source, + component_operator_second_preconditioned_residual, + component_initial, + component_solution, + component_residual, + component_shadow, + component_direction, + component_operator_direction, + residual_squared, + ) + gpu_zero_scalar_accumulator(rho) + gpu_bicgstab_dot_vector(n_cells, component_shadow, component_residual, rho) + gpu_bicgstab_update_direction_vector( + n_cells, + 0.0, + 1.0, + component_residual, + component_direction, + component_operator_direction, + ) + if use_dilu_levels: + gpu_dilu_apply_vector_asymmetric_levels( + dilu_schedule, + owner, + neighbour, + upper, + lower, + preconditioner_reciprocal_diag, + component_direction, + component_preconditioned_direction, + ) + else: + gpu_bicgstab_precondition_vector(n_cells, precond_diag, component_direction, component_preconditioned_direction) + gpu_ldu_matvec_vector_asymmetric_faces( + n_cells, + n_internal_faces, + owner, + neighbour, + upper, + lower, + diag, + component_preconditioned_direction, + component_operator_direction, + ) + gpu_zero_scalar_accumulator(denominator) + gpu_bicgstab_dot_vector(n_cells, component_shadow, component_operator_direction, denominator) + qd.sync() + rho_value = float(np.asarray(rho.to_numpy())[0]) + denominator_value = float(np.asarray(denominator.to_numpy())[0]) + alpha = rho_value/denominator_value + + gpu_zero_scalar_accumulator(residual_squared) + gpu_bicgstab_update_intermediate_vector_preconditioned( + n_cells, + alpha, + component_solution, + component_preconditioned_direction, + component_residual, + component_operator_direction, + component_intermediate_residual, + residual_squared, + ) + qd.sync() + intermediate_residual_squared = float(np.asarray(residual_squared.to_numpy())[0]) + if use_dilu_levels: + gpu_dilu_apply_vector_asymmetric_levels( + dilu_schedule, + owner, + neighbour, + upper, + lower, + preconditioner_reciprocal_diag, + component_intermediate_residual, + component_second_preconditioned_residual, + ) + else: + gpu_bicgstab_precondition_vector(n_cells, precond_diag, component_intermediate_residual, component_second_preconditioned_residual) + gpu_ldu_matvec_vector_asymmetric_faces( + n_cells, + n_internal_faces, + owner, + neighbour, + upper, + lower, + diag, + component_second_preconditioned_residual, + component_operator_second_preconditioned_residual, + ) + gpu_zero_scalar_accumulator(omega_numerator) + gpu_zero_scalar_accumulator(omega_denominator) + gpu_bicgstab_dot_vector(n_cells, component_operator_second_preconditioned_residual, component_intermediate_residual, omega_numerator) + gpu_bicgstab_dot_vector(n_cells, component_operator_second_preconditioned_residual, component_operator_second_preconditioned_residual, omega_denominator) + qd.sync() + omega_numerator_value = float(np.asarray(omega_numerator.to_numpy())[0]) + omega_denominator_value = float(np.asarray(omega_denominator.to_numpy())[0]) + omega = omega_numerator_value/omega_denominator_value + + gpu_zero_scalar_accumulator(residual_squared) + gpu_bicgstab_update_solution_residual_vector_preconditioned( + n_cells, + omega, + component_solution, + component_second_preconditioned_residual, + component_intermediate_residual, + component_operator_second_preconditioned_residual, + component_residual, + residual_squared, + ) + qd.sync() + residual_after_omega_squared = float(np.asarray(residual_squared.to_numpy())[0]) + gpu_scatter_vector_component(n_cells, component, component_intermediate_residual, intermediate_residual_out) + gpu_scatter_vector_component(n_cells, component, component_second_preconditioned_residual, second_preconditioned_residual_out) + gpu_scatter_vector_component(n_cells, component, component_operator_second_preconditioned_residual, operator_second_preconditioned_residual_out) + gpu_scatter_vector_component(n_cells, component, component_residual, residual_after_omega_out) + gpu_scatter_vector_component(n_cells, component, component_solution, solution_after_omega_out) + + rho_values[component] = rho_value + denominator_values[component] = denominator_value + alpha_values[component] = alpha + omega_numerator_values[component] = omega_numerator_value + omega_denominator_values[component] = omega_denominator_value + omega_values[component] = omega + intermediate_residual_squared_values[component] = intermediate_residual_squared + residual_after_omega_squared_values[component] = residual_after_omega_squared + + qd.sync() + return { + "rho": rho_values, + "denominator": denominator_values, + "alpha": alpha_values, + "intermediate_residual_squared": intermediate_residual_squared_values, + "omega_numerator": omega_numerator_values, + "omega_denominator": omega_denominator_values, + "omega": omega_values, + "residual_after_omega_squared": residual_after_omega_squared_values, + "preconditioner": "dilu_level_scheduled" if use_dilu_levels else "diagonal_jacobi", + } + + def gpu_ldu_pbicgstab_vector_asymmetric_components( n_cells: int, n_internal_faces: int, @@ -1741,6 +1990,8 @@ def gpu_ldu_pbicgstab_vector_asymmetric_components( "component_initial_normalized_residual": [report.get("initial_normalized_residual") for report in component_reports], "component_final_normalized_residual": [report.get("final_normalized_residual") for report in component_reports], "component_normalization_factor": [report.get("normalization_factor") for report in component_reports], + "component_iteration_trace": [report.get("iteration_trace", []) for report in component_reports], + "trace_limit": component_reports[0].get("trace_limit") if component_reports else None, "min_iterations": min_iterations, "preconditioner": component_reports[0].get("preconditioner") if component_reports else None, "preconditioner_schedule": component_reports[0].get("preconditioner_schedule") if component_reports else None, diff --git a/scripts/update_loop_diagnostic_context.py b/scripts/update_loop_diagnostic_context.py index a8ac3c8..0423023 100755 --- a/scripts/update_loop_diagnostic_context.py +++ b/scripts/update_loop_diagnostic_context.py @@ -37,7 +37,7 @@ def load_json(path: Path) -> Any | None: def is_verifier_report(value: Any) -> bool: if not isinstance(value, Mapping): return False - return any(key in value for key in ("verifier_evidence", "first_divergence_summary", "artifact_comparisons")) + return any(key in value for key in ("verifier_evidence", "first_divergence_summary", "artifact_comparisons", "differential_trace")) def is_gpu_report(value: Mapping[str, Any], path: Path) -> bool: @@ -175,6 +175,34 @@ def field_checks(report: Mapping[str, Any]) -> list[dict[str, Any]]: out.sort(key=lambda item: (item.get("allclose") is True, item["mode"], str(item.get("field")))) return out +def differential_trace_checks(report: Mapping[str, Any]) -> list[dict[str, Any]]: + trace = report.get("differential_trace", {}) if isinstance(report.get("differential_trace"), Mapping) else {} + checks = trace.get("checkpoints", []) if isinstance(trace.get("checkpoints"), list) else [] + out: list[dict[str, Any]] = [] + for check in checks: + if not isinstance(check, Mapping): + continue + largest = check.get("largest_difference") if isinstance(check.get("largest_difference"), Mapping) else {} + location = largest.get("location") if isinstance(largest.get("location"), Mapping) else {} + metadata = check.get("metadata") if isinstance(check.get("metadata"), Mapping) else {} + diff = check.get("difference_stats") if isinstance(check.get("difference_stats"), Mapping) else {} + out.append( + { + "name": check.get("name"), + "status": check.get("status"), + "reason": check.get("reason"), + "family": metadata.get("family"), + "lifecycle_phase": metadata.get("lifecycle_phase"), + "substitution_supported": metadata.get("substitution_supported"), + "max_abs": diff.get("max_abs"), + "mean_abs": diff.get("mean_abs"), + "rms_abs": diff.get("rms_abs"), + "location": location, + } + ) + out.sort(key=lambda item: (item.get("status") == "passed", str(item.get("name")))) + return out + def extract_metrics(report: Mapping[str, Any]) -> dict[str, Any]: metrics: dict[str, Any] = { @@ -186,6 +214,18 @@ def extract_metrics(report: Mapping[str, Any]) -> dict[str, Any]: metrics["first.target"] = first.get("first_target") metrics["first.family"] = first.get("artifact_family") metrics["first.stage_group"] = first.get("stage_group") + trace = report.get("differential_trace") if isinstance(report.get("differential_trace"), Mapping) else {} + if trace: + first_trace = trace.get("first_divergence") if isinstance(trace.get("first_divergence"), Mapping) else {} + metrics["trace.status"] = trace.get("status") + metrics["trace.first_checkpoint"] = first_trace.get("name") if first_trace else None + metrics["trace.failed_count"] = finite_number(trace.get("failed_count")) or 0 + metrics["trace.missing_count"] = finite_number(trace.get("missing_count")) or 0 + if first_trace: + for key in ("max_abs", "mean_abs", "rms_abs"): + number = finite_number(first_trace.get(key)) + if number is not None: + metrics[f"trace.first.{key}"] = number for check in artifact_checks(report): base = f"artifact.{check['family']}.{check['name']}" metrics[f"{base}.allclose"] = check.get("allclose") is True @@ -293,6 +333,10 @@ def render_context(report_path: Path | None, report: Mapping[str, Any] | None, b deltas = classify_delta(metrics, baseline) checks = artifact_checks(report) fields = field_checks(report) + trace = report.get("differential_trace", {}) if isinstance(report.get("differential_trace"), Mapping) else {} + trace_first = trace.get("first_divergence") if isinstance(trace.get("first_divergence"), Mapping) else {} + trace_substitution = trace.get("substitution") if isinstance(trace.get("substitution"), Mapping) else {} + trace_checks = differential_trace_checks(report) preconditioner = first_nested_key(report, "preconditioner_diagnostic") solver_trace = first_nested_key(report, "linear_solver_trace") @@ -312,11 +356,31 @@ def render_context(report_path: Path | None, report: Mapping[str, Any] | None, b f"- Evidence path: `{first.get('evidence_path') if first else None}`", f"- Field/reason: `{first.get('field') if first else None}` / `{first.get('reason') if first else None}`", "", - "## Solver phase evidence", + "## Differential trace", "", - "| Family | Status | Why |", - "|---|---:|---|", + f"- Status: `{trace.get('status') if trace else None}`", + f"- First checkpoint: `{trace_first.get('name') if trace_first else None}`", + f"- Lifecycle: `{get_path(trace_first, 'metadata.lifecycle_phase') if trace_first else None}`", + f"- max_abs / rms_abs: `{fmt(trace_first.get('max_abs') if trace_first else None)}` / `{fmt(trace_first.get('rms_abs') if trace_first else None)}`", + f"- Substitution requested/loaded: `{trace_substitution.get('requested') if trace_substitution else []}` / `{trace_substitution.get('loaded') if trace_substitution else []}`", + f"- Source artifacts: reference=`{get_path(trace, 'roles.reference.path') if trace else None}`, candidate=`{get_path(trace, 'roles.candidate.path') if trace else None}`", + "", + "| Checkpoint | Status | Lifecycle | max_abs | rms_abs | Location | Substitute? |", + "|---|---:|---|---:|---:|---|---:|", ] + for check in trace_checks[:12]: + lines.append( + "| {name} | {status} | {phase} | {max_abs} | {rms_abs} | {location} | {substitute} |".format( + name=check.get("name"), + status=check.get("status"), + phase=check.get("lifecycle_phase"), + max_abs=fmt(check.get("max_abs")), + rms_abs=fmt(check.get("rms_abs")), + location=render_location(check.get("location")), + substitute=status_word(check.get("substitution_supported")), + ) + ) + lines.extend(["", "## Solver phase evidence", "", "| Family | Status | Why |", "|---|---:|---|"]) for family in families: if not isinstance(family, Mapping): continue @@ -417,6 +481,7 @@ def main(argv: list[str] | None = None) -> int: "report": str(report_path), "metrics": extract_metrics(report), "first_divergence_summary": report.get("first_divergence_summary"), + "differential_trace": report.get("differential_trace"), } args.baseline.parent.mkdir(parents=True, exist_ok=True) args.baseline.write_text(json.dumps(current, indent=2, sort_keys=True) + "\n", encoding="utf-8") diff --git a/scripts/verify_airfrans_stepper.py b/scripts/verify_airfrans_stepper.py index 24dc80f..1ac8129 100755 --- a/scripts/verify_airfrans_stepper.py +++ b/scripts/verify_airfrans_stepper.py @@ -876,20 +876,52 @@ def add_numeric_artifact(arrays: dict[str, np.ndarray], name: str, value: Any) - def add_field_numeric_artifact(arrays: dict[str, np.ndarray], name: str, field: Any) -> None: - if field is not None and hasattr(field, "internal"): + if field is None: + return + if isinstance(field, Mapping): + add_numeric_artifact(arrays, name, field.get("internal")) + elif hasattr(field, "internal"): add_numeric_artifact(arrays, name, getattr(field, "internal")) +def _value_from_mapping_or_attr(value: Any, name: str) -> Any: + if isinstance(value, Mapping): + return value.get(name) + return getattr(value, name, None) + + def add_matrix_numeric_artifacts(arrays: dict[str, np.ndarray], prefix: str, matrix: Any) -> None: if matrix is None: return for component in ("diag", "upper", "lower", "source"): - value = getattr(matrix, component, None) + value = _value_from_mapping_or_attr(matrix, component) if value is not None: add_numeric_artifact(arrays, f"{prefix}.{component}", value) - psi = getattr(matrix, "psi", None) - if psi is not None and hasattr(psi, "internal"): - add_numeric_artifact(arrays, f"{prefix}.psi", getattr(psi, "internal")) + psi = _value_from_mapping_or_attr(matrix, "psi") + if psi is not None: + add_field_numeric_artifact(arrays, f"{prefix}.psi", psi) + + +def add_momentum_term_source_artifacts(arrays: dict[str, np.ndarray], momentum_terms: Any) -> None: + term_outputs = stage_output(momentum_terms, "terms") + if term_outputs is None: + return + source_names = { + "assemble_momentum_ddt": "matrix_terms.UEqn.ddt.source", + "assemble_momentum_div_phi_U": "matrix_terms.UEqn.div.source", + "assemble_momentum_divDevSigma": "matrix_terms.UEqn.divDevSigma.source", + "assemble_momentum_sources": "matrix_terms.UEqn.fvModels.source", + } + for term in term_outputs: + term_name = _value_from_mapping_or_attr(term, "name") + if term_name in source_names: + matrix = _value_from_mapping_or_attr(term, "matrix") + add_numeric_artifact(arrays, source_names[term_name], _value_from_mapping_or_attr(matrix, "source")) + if term_name == "assemble_momentum_divDevSigma": + diagnostics = _value_from_mapping_or_attr(term, "diagnostics") + add_numeric_artifact(arrays, "matrix_terms.UEqn.divDevSigma.wall_dev_tau.source", _value_from_mapping_or_attr(diagnostics, "wall_dev_tau_source")) + elif term_name == "assemble_momentum_MRF_DDt": + add_field_numeric_artifact(arrays, "matrix_terms.UEqn.MRF_DDt.field", _value_from_mapping_or_attr(term, "field")) def write_numeric_artifact_file(path: Path, arrays: Mapping[str, np.ndarray], *, role: str) -> dict[str, Any]: @@ -917,6 +949,8 @@ def write_split_numeric_artifacts( path: Path | None, *, role: str, + momentum_terms: Any | None, + assemble_UEqn: Any, relax_UEqn: Any, solve_UEqn: Any, compute_pressure_inputs: Any, @@ -930,6 +964,15 @@ def write_split_numeric_artifacts( return None arrays: dict[str, np.ndarray] = {} + add_momentum_term_source_artifacts(arrays, momentum_terms) + unrelaxed_UEqn = stage_output(assemble_UEqn, "UEqn") + if unrelaxed_UEqn is not None: + diag = getattr(unrelaxed_UEqn, "diag", None) + if diag is not None: + add_numeric_artifact(arrays, "matrix_operator.UEqn.unrelaxed_diag", diag) + source = getattr(unrelaxed_UEqn, "source", None) + if source is not None: + add_numeric_artifact(arrays, "matrix_operator.UEqn.unrelaxed_source", source) add_matrix_numeric_artifacts(arrays, "matrix_operator.UEqn", stage_output(relax_UEqn, "UEqn")) add_matrix_numeric_artifacts(arrays, "matrix_operator.pEqn", stage_output(assemble_pEqn, "pEqn")) @@ -939,6 +982,23 @@ def write_split_numeric_artifacts( add_field_numeric_artifact(arrays, "solver.solve_UEqn.field_before", stage_output(solve_UEqn, "field_before")) add_field_numeric_artifact(arrays, "solver.solve_UEqn.field_after", stage_output(solve_UEqn, "field_after")) add_matrix_numeric_artifacts(arrays, "solver.solve_UEqn.matrix_before", stage_output(solve_UEqn, "matrix_before")) + add_numeric_artifact(arrays, "solver.solve_UEqn.solve_diag", stage_output(solve_UEqn, "solve_diag")) + add_numeric_artifact(arrays, "solver.solve_UEqn.solve_source", stage_output(solve_UEqn, "solve_source")) + add_numeric_artifact(arrays, "solver.solve_UEqn.initial_residual", stage_output(solve_UEqn, "initial_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.preconditioned_residual", stage_output(solve_UEqn, "preconditioned_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.level_preconditioned_residual", stage_output(solve_UEqn, "level_preconditioned_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.operator_preconditioned_direction", stage_output(solve_UEqn, "operator_preconditioned_direction")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.rho", stage_output(solve_UEqn, "first_iteration_rho")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.denominator", stage_output(solve_UEqn, "first_iteration_denominator")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.alpha", stage_output(solve_UEqn, "first_iteration_alpha")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.intermediate_residual", stage_output(solve_UEqn, "first_iteration_intermediate_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.second_preconditioned_residual", stage_output(solve_UEqn, "first_iteration_second_preconditioned_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual", stage_output(solve_UEqn, "first_iteration_operator_second_preconditioned_residual")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.omega_numerator", stage_output(solve_UEqn, "first_iteration_omega_numerator")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.omega_denominator", stage_output(solve_UEqn, "first_iteration_omega_denominator")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.omega", stage_output(solve_UEqn, "first_iteration_omega")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.residual_after_omega", stage_output(solve_UEqn, "first_iteration_residual_after_omega")) + add_numeric_artifact(arrays, "solver.solve_UEqn.first_iteration.solution_after_omega", stage_output(solve_UEqn, "first_iteration_solution_after_omega")) add_field_numeric_artifact(arrays, "solver.solve_pEqn.field_before", stage_output(solve_pEqn, "field_before")) add_field_numeric_artifact(arrays, "solver.solve_pEqn.p", stage_output(solve_pEqn, "p")) add_field_numeric_artifact(arrays, "solver.solve_pEqn.phi", stage_output(solve_pEqn, "phi")) @@ -970,6 +1030,55 @@ def diagnostic_artifact_path(report: Mapping[str, Any], role: str) -> Path | Non return None return Path(str(value)) +def build_differential_trace(report: Mapping[str, Any], *, mesh_context: Any | None = None) -> dict[str, Any]: + from foam_stepper.differential_trace import build_differential_trace_report, load_npz_artifacts + + reference_path = diagnostic_artifact_path(report, "reference_split") + candidate_path = diagnostic_artifact_path(report, "split") + reference_artifacts = load_npz_artifacts(reference_path) + candidate_artifacts = load_npz_artifacts(candidate_path) + tolerances = report.get("tolerances", {}) if isinstance(report.get("tolerances"), Mapping) else {} + rtol = float(tolerances.get("rtol", 0.0) or 0.0) + atol = float(tolerances.get("atol", 0.0) or 0.0) + + def context_for(name: str, index: tuple[int, ...], shape: tuple[int, ...]) -> dict[str, Any] | None: + location = artifact_location(name, index, shape, mesh_context) + return local_entity_context(mesh_context, location, name) + + substitution = report.get("differential_trace", {}).get("substitution", {}) if isinstance(report.get("differential_trace"), Mapping) else {} + return build_differential_trace_report( + reference_artifacts=reference_artifacts, + candidate_artifacts=candidate_artifacts, + reference_path=reference_path, + candidate_path=candidate_path, + rtol=rtol, + atol=atol, + top_n=NUMERIC_ARTIFACT_TOP_N, + location_context=context_for, + substitution=substitution, + ) + + +def load_trace_substitutions(report: Mapping[str, Any], checkpoints: Iterable[str]) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + from foam_stepper.differential_trace import extract_reference_substitutions, load_npz_artifacts + + reference_path = diagnostic_artifact_path(report, "reference_split") + reference_artifacts = load_npz_artifacts(reference_path) + if reference_artifacts is None: + requested = list(dict.fromkeys(str(item) for item in checkpoints)) + return {}, { + "schema_version": 1, + "status": "failed" if requested else "not_requested", + "requested": requested, + "loaded": [], + "missing_reference": requested, + "unsupported": [], + "source_artifact": str(reference_path) if reference_path is not None else None, + } + substitutions, manifest = extract_reference_substitutions(reference_artifacts, checkpoints) + manifest["source_artifact"] = str(reference_path) if reference_path is not None else None + return substitutions, manifest + def artifact_entity_kind(mesh: Any | None, shape: tuple[int, ...]) -> str | None: if mesh is None or not shape: @@ -1663,8 +1772,8 @@ def prepare_gpu_solver_inputs(foam: Any, stepper: Any, case: Path, prepared: Map return _call_gpu_backend("prepare_gpu_solver_inputs", foam, stepper, case, prepared, backend) -def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None) -> dict[str, Any]: - return _call_gpu_backend("run_gpu_solver_stage_smoke", stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path) +def run_gpu_solver_stage_smoke(stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None, trace_substitutions: Mapping[str, np.ndarray] | None = None) -> dict[str, Any]: + return _call_gpu_backend("run_gpu_solver_stage_smoke", stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path, trace_substitutions=trace_substitutions) def gpu_solver_stage_report(gpu_run: Mapping[str, Any]) -> dict[str, Any]: @@ -1750,8 +1859,8 @@ def run_backend_iteration(stepper: Any, backend: Mapping[str, Any], case: Path) } -def run_gpu_split_iteration(foam: Any, stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None) -> dict[str, Any]: - return _call_gpu_backend("run_gpu_split_iteration", foam, stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path) +def run_gpu_split_iteration(foam: Any, stepper: Any, backend: Mapping[str, Any], case: Path, *, diagnostic_artifact_path: Path | None = None, trace_substitutions: Mapping[str, np.ndarray] | None = None) -> dict[str, Any]: + return _call_gpu_backend("run_gpu_split_iteration", foam, stepper, backend, case, diagnostic_artifact_path=diagnostic_artifact_path, trace_substitutions=trace_substitutions) def make_stepper(foam: Any, case: Path, label: str) -> Any: @@ -1890,6 +1999,8 @@ def run_split_iteration(foam: Any, stepper: Any, *, diagnostic_artifact_path: Pa diagnostic_artifacts = write_split_numeric_artifacts( diagnostic_artifact_path, role="split", + momentum_terms=terms, + assemble_UEqn=UEqn, relax_UEqn=relax_UEqn, solve_UEqn=solve_UEqn, compute_pressure_inputs=compute_pressure_inputs, @@ -2313,6 +2424,34 @@ def first_divergence_summary(report: Mapping[str, Any]) -> dict[str, Any]: for item in diagnostics.get("downstream_unreliable", []) if isinstance(item, Mapping) ] + trace = report.get("differential_trace", {}) if isinstance(report.get("differential_trace"), Mapping) else {} + trace_first = trace.get("first_divergence") if isinstance(trace.get("first_divergence"), Mapping) else None + if trace_first is not None: + metadata = trace_first.get("metadata", {}) if isinstance(trace_first.get("metadata"), Mapping) else {} + largest = trace_first.get("largest_difference", {}) if isinstance(trace_first.get("largest_difference"), Mapping) else {} + return { + "schema_version": 1, + "status": "failed", + "category": COMPARISON_FAILURE, + "first_target": trace_first.get("name"), + "target_kind": "differential_trace_checkpoint", + "evidence_path": "differential_trace.first_divergence", + "artifact_family": metadata.get("family"), + "mode": "split", + "stage_group": metadata.get("family"), + "stage_names": [metadata.get("lifecycle_phase")], + "field": metadata.get("field") or metadata.get("field_or_matrix"), + "reason": trace_first.get("reason"), + "local_context": largest.get("local_context"), + "missing_evidence": None, + "trace_lifecycle_phase": metadata.get("lifecycle_phase"), + "trace_max_abs": trace_first.get("max_abs"), + "trace_mean_abs": trace_first.get("mean_abs"), + "trace_rms_abs": trace_first.get("rms_abs"), + "trace_substitution_supported": metadata.get("substitution_supported"), + "intermediate_artifact_statuses": family_statuses, + "downstream_symptoms_to_ignore": downstream_ignore, + } if first_blocker is not None: field_context = first_blocker.get("field_context", {}) if isinstance(first_blocker.get("field_context"), Mapping) else {} first_target = first_blocker.get("first_target", {}) if isinstance(first_blocker.get("first_target"), Mapping) else {} @@ -2569,6 +2708,8 @@ def build_verifier_evidence(report: Mapping[str, Any]) -> dict[str, Any]: "differentiability_status": report.get("differentiability", {}).get("status"), "mesh_topology_sha256": oracle_mesh.get("topology_sha256"), "first_divergence_summary": report.get("first_divergence_summary"), + "differential_trace_status": report.get("differential_trace", {}).get("status") if isinstance(report.get("differential_trace"), Mapping) else None, + "differential_trace_first_divergence": report.get("differential_trace", {}).get("first_divergence") if isinstance(report.get("differential_trace"), Mapping) else None, "mesh_geometry_sha256": oracle_mesh.get("geometry_sha256"), "modes": regression_modes, }, @@ -2603,6 +2744,8 @@ def base_report(args: argparse.Namespace) -> dict[str, Any]: "skip_split": args.skip_split, "rtol": args.rtol, "atol": args.atol, + "trace_substitute": list(args.trace_substitute), + "trace_substituted_artifact": args.trace_substituted_artifact, }, "prepares_case": True, "runs_oracle": True, @@ -2708,6 +2851,16 @@ def base_report(args: argparse.Namespace) -> dict[str, Any]: "status": "not_evaluated", "root": diagnostic_artifact_root(args), }, + "differential_trace": { + "schema_version": 1, + "status": "not_evaluated", + "comparison_basis": "reference_split_vs_split_npz_checkpoints", + "substitution": { + "requested": list(args.trace_substitute), + "status": "not_requested" if not args.trace_substitute else "pending_reference_artifact", + "cpu_fallback_allowed": False, + }, + }, "gpu_solver_inputs": { "schema_version": GPU_INPUT_SCHEMA_VERSION, "status": "not_requested", @@ -2755,6 +2908,7 @@ def run_harness(args: argparse.Namespace, report: dict[str, Any]) -> None: raise exc.to_harness_error() from exc report["backend"] = json_ready(backend) report["differentiability"] = json_ready(differentiability_report(backend, case_name=args.source.name)) + trace_substitutions: dict[str, np.ndarray] | None = None if backend.get("selected") == "gpu" and split_case is not None: reference_summary, reference_report = run_reference_split_diagnostic(args) report["reference_split_diagnostic"] = json_ready(reference_summary) @@ -2773,6 +2927,14 @@ def run_harness(args: argparse.Namespace, report: dict[str, Any]) -> None: reference_artifacts = reference_report.get("diagnostic_artifacts", {}) if isinstance(reference_artifacts, Mapping) and isinstance(reference_artifacts.get("split"), Mapping): report["diagnostic_artifacts"]["reference_split"] = reference_artifacts["split"] + if args.trace_substitute: + trace_substitutions, trace_substitution_manifest = load_trace_substitutions(report, args.trace_substitute) + report["differential_trace"]["substitution"] = { + **report["differential_trace"].get("substitution", {}), + **trace_substitution_manifest, + "cpu_fallback_allowed": False, + "mode": "reference_checkpoint_replay_before_downstream_gpu_execution", + } if backend.get("selected") == "gpu": try: foam = import_foam() @@ -2920,7 +3082,7 @@ def run_harness(args: argparse.Namespace, report: dict[str, Any]) -> None: split_artifact_path = diagnostic_artifact_root(args) / "split.npz" split_artifact_mesh = split_stepper.mesh() if backend.get("selected") == "gpu": - split_result = run_gpu_split_iteration(foam, split_stepper, backend, split_case, diagnostic_artifact_path=split_artifact_path) + split_result = run_gpu_split_iteration(foam, split_stepper, backend, split_case, diagnostic_artifact_path=split_artifact_path, trace_substitutions=trace_substitutions) else: split_result = run_split_iteration(foam, split_stepper, diagnostic_artifact_path=split_artifact_path) report["stage_observability"]["split"] = split_result["observability"] @@ -2963,6 +3125,30 @@ def run_harness(args: argparse.Namespace, report: dict[str, Any]) -> None: report["comparison_attribution"] = json_ready(comparison_attribution) report["comparison_mismatches"] = json_ready(mismatches) report["artifact_comparisons"] = json_ready(build_artifact_comparisons(report, mesh_context=split_artifact_mesh)) + report["differential_trace"] = build_differential_trace(report, mesh_context=split_artifact_mesh) + if args.trace_substituted_artifact is not None and args.trace_substitute: + from foam_stepper.differential_trace import load_npz_artifacts, materialize_substituted_artifacts + + reference_path = diagnostic_artifact_path(report, "reference_split") + candidate_path = diagnostic_artifact_path(report, "split") + reference_artifacts = load_npz_artifacts(reference_path) + candidate_artifacts = load_npz_artifacts(candidate_path) + if reference_artifacts is not None and candidate_artifacts is not None: + report["differential_trace"]["substituted_artifact"] = materialize_substituted_artifacts( + reference_artifacts=reference_artifacts, + candidate_artifacts=candidate_artifacts, + checkpoints=args.trace_substitute, + output_path=args.trace_substituted_artifact, + ) + else: + report["differential_trace"]["substituted_artifact"] = { + "schema_version": 1, + "status": "failed", + "requested": list(args.trace_substitute), + "reason": "reference_or_candidate_artifact_missing", + "reference_path": str(reference_path) if reference_path is not None else None, + "candidate_path": str(candidate_path) if candidate_path is not None else None, + } report["failure_diagnostics"] = json_ready(build_failure_diagnostics(mismatches, report)) report["timing"] = json_ready(build_timing_evidence(report)) report["intermediate_artifacts"] = json_ready(intermediate_artifact_family_report(report)) @@ -2979,6 +3165,7 @@ def run_harness(args: argparse.Namespace, report: dict[str, Any]) -> None: "likely_causes": comparison_attribution, "first_divergence_summary": report["first_divergence_summary"], "diagnostics": report["failure_diagnostics"], + "differential_trace": report["differential_trace"], }, ) @@ -3043,6 +3230,14 @@ def print_human_summary(report: Mapping[str, Any]) -> None: print(f"intermediate_artifact_families={[(family.get('name'), family.get('status')) for family in families]}") first_family = artifacts.get("first_actionable_family") or {} print(f"intermediate_first_actionable_family={first_family.get('name')}") + trace = report.get("differential_trace") or {} + if trace: + first_trace = trace.get("first_divergence") or {} + substitution = trace.get("substitution") or {} + largest = first_trace.get("largest_difference") or {} + print(f"differential_trace_status={trace.get('status')} checkpoint_count={trace.get('checkpoint_count')} failed={trace.get('failed_count')} missing={trace.get('missing_count')}") + print(f"differential_trace_first_checkpoint={first_trace.get('name')} max_abs={fmt_sci(first_trace.get('max_abs'))} location={format_location((largest.get('location') or {}) if isinstance(largest, Mapping) else {})}") + print(f"differential_trace_substitution_status={substitution.get('status')} loaded={substitution.get('loaded')}") compact = report.get("first_divergence_summary") or {} if compact: print(f"first_divergence_summary={compact}") @@ -3154,6 +3349,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--atol", type=float, default=1e-8) parser.add_argument("--backend", choices=BACKEND_CHOICES, default="auto", help="Backend request; gpu means the full RANS solver backend and currently fails until that backend exists") parser.add_argument("--skip-split", action="store_true", help="Skip the explicit split-step parity check for local debugging") + parser.add_argument("--trace-substitute", action="append", default=[], help="Replace this GPU split checkpoint with the reference_split artifact before downstream replay; may be repeated") + parser.add_argument("--trace-substituted-artifact", type=Path, default=None, help="Optional NPZ path for an artifact-only candidate copy with requested trace substitutions applied") return parser.parse_args(argv) @@ -3165,6 +3362,8 @@ def main(argv: list[str] | None = None) -> int: args.report = args.work / "verifier_report.json" else: args.report = args.report.resolve() + if args.trace_substituted_artifact is not None: + args.trace_substituted_artifact = args.trace_substituted_artifact.resolve() report = base_report(args) exit_code = 0