From 83e0e8294b56532462925b0f62158517939cd213 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 29 May 2026 14:51:12 -0700 Subject: [PATCH 01/10] MGN diffusion test. --- tests/AMSlib/ams_interface/CMakeLists.txt | 5 + .../test_graph_mgn_surrogate.cpp | 105 ++++ tests/AMSlib/models/CMakeLists.txt | 48 ++ .../models/generate_mgn_graph_diffusion.py | 567 ++++++++++++++++++ 4 files changed, 725 insertions(+) create mode 100644 tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp create mode 100644 tests/AMSlib/models/generate_mgn_graph_diffusion.py diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index 17ee7034..d8d2669f 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -33,3 +33,8 @@ ADD_AMS_UNIT_TEST(AMS_GRAPH_FALLBACK ams_graph_fallback) BUILD_UNIT_TEST(ams_graph_surrogate test_graph_surrogate.cpp) ADD_AMS_UNIT_TEST(AMS_GRAPH_SURROGATE ams_graph_surrogate) +BUILD_UNIT_TEST(ams_graph_mgn_surrogate test_graph_mgn_surrogate.cpp) +if (TARGET generate_mgn_graph_diffusion) + add_dependencies(ams_graph_mgn_surrogate generate_mgn_graph_diffusion) +endif() +ADD_AMS_UNIT_TEST(AMS_GRAPH_MGN_SURROGATE ams_graph_mgn_surrogate) diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp new file mode 100644 index 00000000..05f9a4c9 --- /dev/null +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -0,0 +1,105 @@ +#include +#include + +#include +#include +#include +#include + +#include "AMS.h" +#include "AMSGraph.hpp" +#include "AMSTensor.hpp" +#include "models/mgn_graph_fixtures.hpp" + +using namespace ams; + +using Dim = AMSTensor::IntDimType; + +static const char* MGN_GRAPH_MODEL_PATH = "../models/mgn_graph_diffusion.pt"; + +static std::vector contiguousStrides(const std::vector& shape) +{ + std::vector strides(shape.size(), 1); + Dim stride = 1; + for (std::size_t i = shape.size(); i-- > 0;) { + strides[i] = stride; + stride *= shape[i]; + } + return strides; +} + +template +static AMSTensor makeTensor(std::vector shape, const T* values) +{ + std::vector strides = contiguousStrides(shape); + auto tensor = AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); + std::copy(values, values + tensor.elements(), tensor.template data()); + return tensor; +} + +static Dim toDim(std::int64_t value) { return static_cast(value); } + +static AMSHomogeneousGraph makeGraph(const ams::test::mgn::GraphFixture& f) +{ + return AMSHomogeneousGraph( + makeTensor({toDim(f.num_nodes), toDim(f.node_feature_dim)}, + f.node_features), + makeTensor({2, toDim(f.num_edges)}, f.edge_index), + makeTensor({toDim(f.num_edges), toDim(f.edge_feature_dim)}, + f.edge_features), + makeTensor({1, toDim(f.global_feature_dim)}, f.global_features)); +} + +static void verifyDeltaU(const ams::test::mgn::GraphFixture& f, + const AMSHomogeneousGraphFields& outputs) +{ + CATCH_REQUIRE(outputs.node_fields.contains("delta_u")); + const auto& delta_u = outputs.node_fields.at("delta_u"); + CATCH_REQUIRE(delta_u.shape()[0] == f.num_nodes); + CATCH_REQUIRE(delta_u.shape()[1] == f.reference_output_dim); + + const float* actual = delta_u.data(); + const std::int64_t count = f.num_nodes * f.reference_output_dim; + for (std::int64_t i = 0; i < count; ++i) { + CATCH_REQUIRE(actual[i] == Catch::Approx(f.reference_delta_u[i]) + .epsilon(ams::test::mgn::kComparisonRtol) + .margin(ams::test::mgn::kComparisonAtol)); + } +} + +CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", + "[wf][graph][surrogate][mgn]") +{ + AMSInit(); + + auto model = AMSRegisterAbstractModel("test_mgn_graph_diffusion_surrogate", + 0.5, + MGN_GRAPH_MODEL_PATH, + false); + AMSExecutor executor = AMSCreateExecutor(model, 0, 1); + + for (std::size_t i = 0; i < ams::test::mgn::kNumGraphFixtures; ++i) { + const auto& fixture = ams::test::mgn::kGraphFixtures[i]; + CATCH_DYNAMIC_SECTION("fixture " << fixture.name) + { + AMSHomogeneousGraph graph = makeGraph(fixture); + + bool callback_invoked = false; + HomogeneousGraphDomainFn callback = + [&](const AMSHomogeneousGraph&, AMSHomogeneousGraphFields& outputs) { + callback_invoked = true; + outputs.node_fields.set( + "delta_u", + makeTensor({toDim(fixture.num_nodes), + toDim(fixture.reference_output_dim)}, + fixture.reference_delta_u)); + }; + + AMSHomogeneousGraphFields outputs; + AMSExecute(executor, callback, graph, outputs); + + CATCH_REQUIRE_FALSE(callback_invoked); + verifyDeltaU(fixture, outputs); + } + } +} diff --git a/tests/AMSlib/models/CMakeLists.txt b/tests/AMSlib/models/CMakeLists.txt index ee1e5847..e3f874f3 100644 --- a/tests/AMSlib/models/CMakeLists.txt +++ b/tests/AMSlib/models/CMakeLists.txt @@ -4,6 +4,7 @@ set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/generate_linear_model.py" "${CMAKE_CURRENT_SOURCE_DIR}/generate_base_models.py" "${CMAKE_CURRENT_SOURCE_DIR}/generate_graph_models.py" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" ) # Where to drop the .pt files @@ -114,3 +115,50 @@ configure_file( ${CMAKE_CURRENT_BINARY_DIR}/ams_test_linear_models.hpp COPYONLY ) + +find_program(AMS_PYTHON_EXECUTABLE NAMES python3 python REQUIRED) + +set(MGN_GRAPH_DIFFUSION_CHECKPOINT + "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion_checkpoint.pt" +) +set(MGN_GRAPH_DIFFUSION_MODEL + "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion.pt" +) +set(MGN_GRAPH_DIFFUSION_FIXTURES + "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_fixtures.hpp" +) + +add_custom_command( + OUTPUT + "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" + "${MGN_GRAPH_DIFFUSION_MODEL}" + "${MGN_GRAPH_DIFFUSION_FIXTURES}" + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --mode feasibility + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --mode train + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --mode fixtures + DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + "${CMAKE_CURRENT_SOURCE_DIR}/ams_model.py" + WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}" + COMMENT "Generating Stage 2 pure-Torch MGN graph diffusion model and fixtures" + VERBATIM +) + +add_custom_target(generate_mgn_graph_diffusion + DEPENDS + "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" + "${MGN_GRAPH_DIFFUSION_MODEL}" + "${MGN_GRAPH_DIFFUSION_FIXTURES}" +) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py new file mode 100644 index 00000000..cda2640c --- /dev/null +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -0,0 +1,567 @@ +#!/usr/bin/env python3 +"""Train/export a tiny pure-Torch MGN-like graph diffusion surrogate. + +This Stage 2 utility intentionally stays pure PyTorch. It generates synthetic +homogeneous graphs, trains a small message-passing model, exports an +AMS-wrapped TorchScript model, and writes C++ fixtures whose reference outputs +come from the exported TorchScript model. +""" + +import argparse +import os +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + +try: + import torch + import torch.nn as nn + from torch import Tensor +except ModuleNotFoundError as exc: + raise SystemExit( + "PyTorch is required for Stage 2 MGN graph diffusion generation. " + "Install PyTorch or run this script on a Torch-enabled machine." + ) from exc + + +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(CURRENT_DIR) +from ams_model import create_ams_homogeneous_graph_model + + +NODE_FEATURE_DIM = 4 +EDGE_FEATURE_DIM = 4 +GLOBAL_FEATURE_DIM = 1 +REFERENCE_OUTPUT_DIM = 1 + +FIXTURE_GRAPH_SIZES = (24, 73) +FEASIBILITY_SEEDS = (70024, 70073) +FIXTURE_SEEDS = (90024, 90073) + +MODEL_SEED = 314159 +TRAIN_DATA_SEED = 271828 +VAL_DATA_SEED = 161803 + +LATENT_DIM = 32 +NUM_PROCESSOR_BLOCKS = 2 +K_NEIGHBORS = 6 +TRAIN_STEPS = 2000 +VAL_GRAPHS = 32 +LEARNING_RATE = 1.0e-3 +WEIGHT_DECAY = 1.0e-6 +PRIMARY_BASELINE_FACTOR = 0.1 +SECONDARY_ABSOLUTE_MSE = 1.0e-4 + +CHECKPOINT_NAME = "mgn_graph_diffusion_checkpoint.pt" +MODEL_NAME = "mgn_graph_diffusion.pt" +FIXTURE_HEADER_NAME = "mgn_graph_fixtures.hpp" + + +class MLP(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.layers(x) + + +class ProcessorBlock(nn.Module): + def __init__(self, latent_dim: int, global_dim: int): + super().__init__() + self.edge_mlp = MLP(latent_dim * 3 + global_dim, latent_dim, latent_dim) + self.node_mlp = MLP(latent_dim * 2 + global_dim, latent_dim, latent_dim) + + def forward( + self, + node_latent: Tensor, + edge_latent: Tensor, + edge_index: Tensor, + global_features: Tensor, + ) -> Tuple[Tensor, Tensor]: + src = edge_index[0] + dst = edge_index[1] + + global_edges = global_features.expand(edge_latent.shape[0], global_features.shape[1]) + edge_input = torch.cat( + ( + edge_latent, + node_latent.index_select(0, src), + node_latent.index_select(0, dst), + global_edges, + ), + dim=1, + ) + edge_latent = edge_latent + self.edge_mlp(edge_input) + + aggregated = torch.zeros( + (node_latent.shape[0], edge_latent.shape[1]), + dtype=node_latent.dtype, + device=node_latent.device, + ) + aggregated = aggregated.index_add(0, dst, edge_latent) + + global_nodes = global_features.expand(node_latent.shape[0], global_features.shape[1]) + node_input = torch.cat((node_latent, aggregated, global_nodes), dim=1) + node_latent = node_latent + self.node_mlp(node_input) + return node_latent, edge_latent + + +class TinyGraphDiffusionMGN(nn.Module): + def __init__( + self, + node_dim: int = NODE_FEATURE_DIM, + edge_dim: int = EDGE_FEATURE_DIM, + global_dim: int = GLOBAL_FEATURE_DIM, + latent_dim: int = LATENT_DIM, + ): + super().__init__() + self.node_encoder = MLP(node_dim, latent_dim, latent_dim) + self.edge_encoder = MLP(edge_dim, latent_dim, latent_dim) + self.processor1 = ProcessorBlock(latent_dim, global_dim) + self.processor2 = ProcessorBlock(latent_dim, global_dim) + self.node_decoder = MLP(latent_dim, latent_dim, REFERENCE_OUTPUT_DIM) + + def forward(self, graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + node_features = graph["node_features"] + edge_index = graph["edge_index"].to(torch.int64) + edge_features = graph["edge_features"] + global_features = graph["global_features"] + + node_latent = self.node_encoder(node_features) + edge_latent = self.edge_encoder(edge_features) + node_latent, edge_latent = self.processor1( + node_latent, edge_latent, edge_index, global_features + ) + node_latent, edge_latent = self.processor2( + node_latent, edge_latent, edge_index, global_features + ) + delta_u = self.node_decoder(node_latent) + return {"node:delta_u": delta_u} + + +def model_config() -> Dict[str, int]: + return { + "node_feature_dim": NODE_FEATURE_DIM, + "edge_feature_dim": EDGE_FEATURE_DIM, + "global_feature_dim": GLOBAL_FEATURE_DIM, + "reference_output_dim": REFERENCE_OUTPUT_DIM, + "latent_dim": LATENT_DIM, + "num_processor_blocks": NUM_PROCESSOR_BLOCKS, + "k_neighbors": K_NEIGHBORS, + } + + +def make_model() -> TinyGraphDiffusionMGN: + torch.manual_seed(MODEL_SEED) + model = TinyGraphDiffusionMGN() + return model.to(device=torch.device("cpu"), dtype=torch.float32) + + +def make_generator(seed: int) -> torch.Generator: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + return generator + + +def _unique_directed_edges(knn: Tensor) -> Tensor: + pairs = set() + num_nodes = int(knn.shape[0]) + for dst in range(num_nodes): + for n in range(int(knn.shape[1])): + src = int(knn[dst, n].item()) + if src == dst: + continue + pairs.add((src, dst)) + pairs.add((dst, src)) + ordered = sorted(pairs) + return torch.tensor(ordered, dtype=torch.int64).t().contiguous() + + +def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor]: + generator = make_generator(seed) + positions = torch.rand((num_nodes, 2), generator=generator, dtype=torch.float32) + u = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) * 2.0 - 1.0 + kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 + dt = torch.rand((1, 1), generator=generator, dtype=torch.float32) * 0.06 + 0.02 + + distances = torch.cdist(positions, positions) + masked = distances + torch.eye(num_nodes, dtype=torch.float32) * 1.0e6 + knn = torch.topk(masked, k=K_NEIGHBORS, largest=False, dim=1).indices + edge_index = _unique_directed_edges(knn) + src = edge_index[0] + dst = edge_index[1] + + delta_pos = positions.index_select(0, src) - positions.index_select(0, dst) + distance = delta_pos.norm(dim=1, keepdim=True).clamp_min(1.0e-6) + conductivity = 0.5 * (kappa.index_select(0, src) + kappa.index_select(0, dst)) + raw_weight = conductivity / (distance + 0.05) + incoming_sum = torch.zeros((num_nodes, 1), dtype=torch.float32) + incoming_sum = incoming_sum.index_add(0, dst, raw_weight) + weight = raw_weight / incoming_sum.index_select(0, dst).clamp_min(1.0e-8) + + node_features = torch.cat((positions, u, kappa), dim=1).contiguous() + edge_features = torch.cat((delta_pos, distance, weight), dim=1).contiguous() + messages = weight * (u.index_select(0, src) - u.index_select(0, dst)) + target = torch.zeros((num_nodes, 1), dtype=torch.float32) + target = target.index_add(0, dst, messages) + target = dt * target + + graph = { + "node_features": node_features, + "edge_index": edge_index, + "edge_features": edge_features, + "global_features": dt.contiguous(), + } + return graph, target.contiguous() + + +def graph_to_model_input(graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + return { + "node_features": graph["node_features"].to(dtype=torch.float32), + "edge_index": graph["edge_index"].to(dtype=torch.int64), + "edge_features": graph["edge_features"].to(dtype=torch.float32), + "global_features": graph["global_features"].to(dtype=torch.float32), + } + + +def assert_close(name: str, actual: Tensor, expected: Tensor, atol: float = 1.0e-6, rtol: float = 1.0e-6) -> None: + if not torch.allclose(actual, expected, atol=atol, rtol=rtol): + max_diff = (actual - expected).abs().max().item() + raise RuntimeError(f"{name} mismatch: max abs diff={max_diff}") + + +def run_feasibility(out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + print("[info] Stage 2 feasibility gate") + print(f"[info] model_seed={MODEL_SEED}, fixture_sizes={FIXTURE_GRAPH_SIZES}, config={model_config()}") + + model = make_model().eval() + graphs = [graph_to_model_input(generate_graph(n, seed)[0]) for n, seed in zip(FIXTURE_GRAPH_SIZES, FEASIBILITY_SEEDS)] + + with torch.no_grad(): + eager = [model(graph)["node:delta_u"].detach() for graph in graphs] + + scripted = torch.jit.script(model) + with torch.no_grad(): + scripted_out = [scripted(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(scripted_out, eager)): + assert_close(f"scripted graph {i}", actual, expected) + + script_path = out_dir / "mgn_graph_diffusion_feasibility.pt" + scripted.save(str(script_path)) + reloaded = torch.jit.load(str(script_path)).eval() + with torch.no_grad(): + reloaded_out = [reloaded(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(reloaded_out, eager)): + assert_close(f"reloaded graph {i}", actual, expected) + + wrapped = create_ams_homogeneous_graph_model(model, torch.device("cpu"), torch.float32) + with torch.no_grad(): + wrapped_out = [wrapped(graph)["node:delta_u"].detach() for graph in graphs] + for i, (actual, expected) in enumerate(zip(wrapped_out, eager)): + assert_close(f"AMS-wrapped graph {i}", actual, expected) + + print(f"[info] Feasibility passed for dynamic graph sizes {FIXTURE_GRAPH_SIZES}") + + +def validation_graphs() -> List[Tuple[Dict[str, Tensor], Tensor]]: + graphs = [] + for i in range(VAL_GRAPHS): + num_nodes = 16 + ((i * 37) % 113) + graphs.append(generate_graph(num_nodes, VAL_DATA_SEED + i)) + return graphs + + +def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]]) -> Tuple[float, float, Tensor]: + total_loss = 0.0 + total_baseline = 0.0 + total_nodes = 0 + targets = [] + model.eval() + with torch.no_grad(): + for graph, target in cases: + pred = model(graph_to_model_input(graph))["node:delta_u"] + loss_sum = (pred - target).pow(2).sum().item() + baseline_sum = target.pow(2).sum().item() + total_loss += loss_sum + total_baseline += baseline_sum + total_nodes += int(target.shape[0]) + targets.append(target) + return total_loss / total_nodes, total_baseline / total_nodes, torch.cat(targets, dim=0) + + +def run_train(out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + print("[info] Stage 2 training") + print(f"[info] model_seed={MODEL_SEED}, train_data_seed={TRAIN_DATA_SEED}, val_data_seed={VAL_DATA_SEED}") + print(f"[info] config={model_config()}") + print( + "[info] training " + f"steps={TRAIN_STEPS}, lr={LEARNING_RATE}, weight_decay={WEIGHT_DECAY}, " + f"primary_pass=validation_mse <= {PRIMARY_BASELINE_FACTOR} * zero_baseline_mse" + ) + + model = make_model().train() + optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) + val_cases = validation_graphs() + + for step in range(1, TRAIN_STEPS + 1): + num_nodes = 16 + ((step * 53) % 113) + graph, target = generate_graph(num_nodes, TRAIN_DATA_SEED + step) + pred = model(graph_to_model_input(graph))["node:delta_u"] + loss = torch.nn.functional.mse_loss(pred, target) + + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + if step == 1 or step % 200 == 0 or step == TRAIN_STEPS: + val_mse, zero_baseline_mse, _ = evaluate(model, val_cases) + print( + f"[info] step={step:04d} train_mse={loss.item():.8e} " + f"val_mse={val_mse:.8e} zero_baseline_mse={zero_baseline_mse:.8e}" + ) + + val_mse, zero_baseline_mse, val_targets = evaluate(model, val_cases) + target_mean = val_targets.mean().item() + target_max_abs = val_targets.abs().max().item() + target_std = val_targets.std(unbiased=False).item() + primary_threshold = PRIMARY_BASELINE_FACTOR * zero_baseline_mse + + print( + "[info] target_delta_u stats " + f"mean={target_mean:.8e} max_abs={target_max_abs:.8e} std={target_std:.8e}" + ) + print( + f"[info] final validation_mse={val_mse:.8e}, " + f"zero_baseline_mse={zero_baseline_mse:.8e}, primary_threshold={primary_threshold:.8e}" + ) + if val_mse > SECONDARY_ABSOLUTE_MSE: + print( + f"[warn] validation_mse={val_mse:.8e} is above secondary diagnostic " + f"absolute MSE {SECONDARY_ABSOLUTE_MSE:.8e}" + ) + if val_mse > primary_threshold: + raise RuntimeError( + "Training failed primary Stage 2 criterion: " + f"validation_mse={val_mse:.8e} > {primary_threshold:.8e}" + ) + + metadata = { + "model_config": model_config(), + "model_seed": MODEL_SEED, + "train_data_seed": TRAIN_DATA_SEED, + "val_data_seed": VAL_DATA_SEED, + "fixture_graph_sizes": list(FIXTURE_GRAPH_SIZES), + "fixture_seeds": list(FIXTURE_SEEDS), + "train_steps": TRAIN_STEPS, + "learning_rate": LEARNING_RATE, + "weight_decay": WEIGHT_DECAY, + "validation_mse": val_mse, + "zero_baseline_mse": zero_baseline_mse, + "primary_baseline_factor": PRIMARY_BASELINE_FACTOR, + "secondary_absolute_mse": SECONDARY_ABSOLUTE_MSE, + "target_mean": target_mean, + "target_max_abs": target_max_abs, + "target_std": target_std, + "torch_version": torch.__version__, + } + checkpoint = { + "model_state_dict": model.state_dict(), + "metadata": metadata, + } + checkpoint_path = out_dir / CHECKPOINT_NAME + torch.save(checkpoint, checkpoint_path) + print(f"[info] Saved checkpoint: {checkpoint_path}") + + +def format_float(value: float) -> str: + text = f"{float(value):.9g}" + if "e" not in text and "." not in text: + text += ".0" + return text + "f" + + +def format_array(name: str, c_type: str, values: List[float], values_per_line: int = 6) -> str: + lines = [f"inline constexpr {c_type} {name}[] = {{"] + for i in range(0, len(values), values_per_line): + chunk = values[i : i + values_per_line] + if c_type == "float": + rendered = ", ".join(format_float(v) for v in chunk) + else: + rendered = ", ".join(str(int(v)) for v in chunk) + lines.append(f" {rendered},") + lines.append("};") + return "\n".join(lines) + + +def flatten_float(tensor: Tensor) -> List[float]: + return [float(x) for x in tensor.detach().cpu().contiguous().view(-1).tolist()] + + +def flatten_int(tensor: Tensor) -> List[int]: + return [int(x) for x in tensor.detach().cpu().contiguous().view(-1).tolist()] + + +def write_fixture_header(path: Path, cases: List[Dict[str, object]], metadata: Dict[str, object]) -> None: + lines = [ + "#pragma once", + "", + "#include ", + "#include ", + "", + "// Generated by tests/AMSlib/models/generate_mgn_graph_diffusion.py --mode fixtures.", + f"// model_seed: {metadata['model_seed']}", + f"// train_data_seed: {metadata['train_data_seed']}", + f"// val_data_seed: {metadata['val_data_seed']}", + f"// fixture_graph_sizes: {metadata['fixture_graph_sizes']}", + f"// fixture_seeds: {metadata['fixture_seeds']}", + f"// model_config: {metadata['model_config']}", + f"// validation_mse: {metadata['validation_mse']:.9e}", + f"// zero_baseline_mse: {metadata['zero_baseline_mse']:.9e}", + f"// primary criterion: validation_mse <= {metadata['primary_baseline_factor']} * zero_baseline_mse", + "", + "namespace ams::test::mgn", + "{", + "", + "struct GraphFixture {", + " const char* name;", + " std::int64_t num_nodes;", + " std::int64_t num_edges;", + " std::int64_t node_feature_dim;", + " std::int64_t edge_feature_dim;", + " std::int64_t global_feature_dim;", + " std::int64_t reference_output_dim;", + " const float* node_features;", + " const std::int64_t* edge_index;", + " const float* edge_features;", + " const float* global_features;", + " const float* reference_delta_u;", + "};", + "", + "inline constexpr float kComparisonRtol = 2.0e-5f;", + "inline constexpr float kComparisonAtol = 2.0e-5f;", + "", + ] + + for case in cases: + prefix = str(case["prefix"]) + lines.append(f"// {case['name']}: N={case['num_nodes']}, E={case['num_edges']}") + lines.append(format_array(f"{prefix}_node_features", "float", case["node_features"])) + lines.append(format_array(f"{prefix}_edge_index", "std::int64_t", case["edge_index"], values_per_line=8)) + lines.append(format_array(f"{prefix}_edge_features", "float", case["edge_features"])) + lines.append(format_array(f"{prefix}_global_features", "float", case["global_features"])) + lines.append(format_array(f"{prefix}_reference_delta_u", "float", case["reference_delta_u"])) + lines.append("") + + lines.append("inline constexpr GraphFixture kGraphFixtures[] = {") + for case in cases: + prefix = str(case["prefix"]) + lines.extend( + [ + " {", + f" \"{case['name']}\",", + f" {case['num_nodes']},", + f" {case['num_edges']},", + f" {NODE_FEATURE_DIM},", + f" {EDGE_FEATURE_DIM},", + f" {GLOBAL_FEATURE_DIM},", + f" {REFERENCE_OUTPUT_DIM},", + f" {prefix}_node_features,", + f" {prefix}_edge_index,", + f" {prefix}_edge_features,", + f" {prefix}_global_features,", + f" {prefix}_reference_delta_u,", + " },", + ] + ) + lines.extend( + [ + "};", + "", + "inline constexpr std::size_t kNumGraphFixtures =", + " sizeof(kGraphFixtures) / sizeof(kGraphFixtures[0]);", + "", + "} // namespace ams::test::mgn", + "", + ] + ) + path.write_text("\n".join(lines), encoding="utf-8") + + +def run_fixtures(out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = out_dir / CHECKPOINT_NAME + if not checkpoint_path.exists(): + raise RuntimeError( + f"Missing checkpoint {checkpoint_path}. Run --mode train before --mode fixtures." + ) + + checkpoint = torch.load(checkpoint_path, map_location="cpu") + metadata = checkpoint["metadata"] + print(f"[info] Loading checkpoint: {checkpoint_path}") + print(f"[info] metadata={metadata}") + + model = make_model().eval() + model.load_state_dict(checkpoint["model_state_dict"]) + wrapped = create_ams_homogeneous_graph_model(model, torch.device("cpu"), torch.float32) + + model_path = out_dir / MODEL_NAME + wrapped.save(str(model_path)) + print(f"[info] Saved AMS TorchScript model: {model_path}") + + reloaded = torch.jit.load(str(model_path)).eval() + cases: List[Dict[str, object]] = [] + with torch.no_grad(): + for index, (num_nodes, seed) in enumerate(zip(FIXTURE_GRAPH_SIZES, FIXTURE_SEEDS)): + graph, _target = generate_graph(num_nodes, seed) + graph = graph_to_model_input(graph) + output = reloaded(graph)["node:delta_u"].detach().cpu().contiguous() + case = { + "name": f"mgn_diffusion_n{num_nodes}", + "prefix": f"kCase{index}", + "num_nodes": int(num_nodes), + "num_edges": int(graph["edge_index"].shape[1]), + "node_features": flatten_float(graph["node_features"]), + "edge_index": flatten_int(graph["edge_index"]), + "edge_features": flatten_float(graph["edge_features"]), + "global_features": flatten_float(graph["global_features"]), + "reference_delta_u": flatten_float(output), + } + print( + f"[info] fixture {case['name']} N={case['num_nodes']} " + f"E={case['num_edges']} reference_shape={tuple(output.shape)}" + ) + cases.append(case) + + header_path = out_dir / FIXTURE_HEADER_NAME + write_fixture_header(header_path, cases, metadata) + print(f"[info] Saved C++ fixture header: {header_path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-dir", type=Path, required=True, help="Output directory for Stage 2 artifacts") + parser.add_argument( + "--mode", + choices=("feasibility", "train", "fixtures", "all"), + required=True, + help="Separated Stage 2 generation mode", + ) + args = parser.parse_args() + + if args.mode in ("feasibility", "all"): + run_feasibility(args.out_dir) + if args.mode in ("train", "all"): + run_train(args.out_dir) + if args.mode in ("fixtures", "all"): + run_fixtures(args.out_dir) + + +if __name__ == "__main__": + main() From d4c70471fae9756b10b8621fecff6b2142219d60 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 29 May 2026 15:11:53 -0700 Subject: [PATCH 02/10] Improve model and training. --- .../models/generate_mgn_graph_diffusion.py | 43 ++++++++++++++++--- 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index cda2640c..28465831 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -8,6 +8,7 @@ """ import argparse +import copy import os import sys from pathlib import Path @@ -42,10 +43,10 @@ TRAIN_DATA_SEED = 271828 VAL_DATA_SEED = 161803 -LATENT_DIM = 32 +LATENT_DIM = 64 NUM_PROCESSOR_BLOCKS = 2 K_NEIGHBORS = 6 -TRAIN_STEPS = 2000 +TRAIN_STEPS = 3000 VAL_GRAPHS = 32 LEARNING_RATE = 1.0e-3 WEIGHT_DECAY = 1.0e-6 @@ -62,7 +63,7 @@ def __init__(self, input_dim: int, hidden_dim: int, output_dim: int): super().__init__() self.layers = nn.Sequential( nn.Linear(input_dim, hidden_dim), - nn.ReLU(), + nn.SiLU(), nn.Linear(hidden_dim, output_dim), ) @@ -185,7 +186,17 @@ def _unique_directed_edges(knn: Tensor) -> Tensor: def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor]: generator = make_generator(seed) positions = torch.rand((num_nodes, 2), generator=generator, dtype=torch.float32) - u = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) * 2.0 - 1.0 + phases = torch.rand((1, 4), generator=generator, dtype=torch.float32) * (2.0 * torch.pi) + amps = torch.rand((1, 4), generator=generator, dtype=torch.float32) * 0.5 + 0.5 + x = positions[:, 0:1] + y = positions[:, 1:2] + u_raw = ( + amps[:, 0:1] * torch.sin(2.0 * torch.pi * x + phases[:, 0:1]) + + amps[:, 1:2] * torch.cos(2.0 * torch.pi * y + phases[:, 1:2]) + + amps[:, 2:3] * torch.sin(2.0 * torch.pi * (x + y) + phases[:, 2:3]) + + amps[:, 3:4] * torch.cos(2.0 * torch.pi * (x - y) + phases[:, 3:4]) + ) + u = torch.tanh(0.5 * u_raw) kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 dt = torch.rand((1, 1), generator=generator, dtype=torch.float32) * 0.06 + 0.02 @@ -308,7 +319,16 @@ def run_train(out_dir: Path) -> None: model = make_model().train() optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) + scheduler = torch.optim.lr_scheduler.MultiStepLR( + optimizer, + milestones=(TRAIN_STEPS // 2, (TRAIN_STEPS * 5) // 6), + gamma=0.3, + ) val_cases = validation_graphs() + best_val_mse = float("inf") + best_zero_baseline_mse = float("inf") + best_step = 0 + best_state_dict = copy.deepcopy(model.state_dict()) for step in range(1, TRAIN_STEPS + 1): num_nodes = 16 + ((step * 53) % 113) @@ -319,14 +339,22 @@ def run_train(out_dir: Path) -> None: optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() + scheduler.step() if step == 1 or step % 200 == 0 or step == TRAIN_STEPS: val_mse, zero_baseline_mse, _ = evaluate(model, val_cases) + if val_mse < best_val_mse: + best_val_mse = val_mse + best_zero_baseline_mse = zero_baseline_mse + best_step = step + best_state_dict = copy.deepcopy(model.state_dict()) print( f"[info] step={step:04d} train_mse={loss.item():.8e} " - f"val_mse={val_mse:.8e} zero_baseline_mse={zero_baseline_mse:.8e}" + f"val_mse={val_mse:.8e} zero_baseline_mse={zero_baseline_mse:.8e} " + f"best_step={best_step:04d} best_val_mse={best_val_mse:.8e}" ) + model.load_state_dict(best_state_dict) val_mse, zero_baseline_mse, val_targets = evaluate(model, val_cases) target_mean = val_targets.mean().item() target_max_abs = val_targets.abs().max().item() @@ -338,7 +366,7 @@ def run_train(out_dir: Path) -> None: f"mean={target_mean:.8e} max_abs={target_max_abs:.8e} std={target_std:.8e}" ) print( - f"[info] final validation_mse={val_mse:.8e}, " + f"[info] selected best_step={best_step}, validation_mse={val_mse:.8e}, " f"zero_baseline_mse={zero_baseline_mse:.8e}, primary_threshold={primary_threshold:.8e}" ) if val_mse > SECONDARY_ABSOLUTE_MSE: @@ -363,6 +391,9 @@ def run_train(out_dir: Path) -> None: "learning_rate": LEARNING_RATE, "weight_decay": WEIGHT_DECAY, "validation_mse": val_mse, + "best_step": best_step, + "best_validation_mse": best_val_mse, + "best_zero_baseline_mse": best_zero_baseline_mse, "zero_baseline_mse": zero_baseline_mse, "primary_baseline_factor": PRIMARY_BASELINE_FACTOR, "secondary_absolute_mse": SECONDARY_ABSOLUTE_MSE, From 8324d3384cb8cec84fcee893e1108821f540c633 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 29 May 2026 15:23:43 -0700 Subject: [PATCH 03/10] REfine model. --- tests/AMSlib/models/generate_mgn_graph_diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 28465831..6d95ba48 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -216,8 +216,8 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor weight = raw_weight / incoming_sum.index_select(0, dst).clamp_min(1.0e-8) node_features = torch.cat((positions, u, kappa), dim=1).contiguous() - edge_features = torch.cat((delta_pos, distance, weight), dim=1).contiguous() messages = weight * (u.index_select(0, src) - u.index_select(0, dst)) + edge_features = torch.cat((delta_pos, distance, messages), dim=1).contiguous() target = torch.zeros((num_nodes, 1), dtype=torch.float32) target = target.index_add(0, dst, messages) target = dt * target From cb6617a3d99474e67c81b4d6a15dbc3a3af28199 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 29 May 2026 15:36:39 -0700 Subject: [PATCH 04/10] load checkpoint fix. --- tests/AMSlib/models/generate_mgn_graph_diffusion.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 6d95ba48..5c05a5bc 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -400,7 +400,7 @@ def run_train(out_dir: Path) -> None: "target_mean": target_mean, "target_max_abs": target_max_abs, "target_std": target_std, - "torch_version": torch.__version__, + "torch_version": str(torch.__version__), } checkpoint = { "model_state_dict": model.state_dict(), @@ -439,6 +439,13 @@ def flatten_int(tensor: Tensor) -> List[int]: return [int(x) for x in tensor.detach().cpu().contiguous().view(-1).tolist()] +def load_checkpoint(path: Path) -> Dict[str, object]: + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + def write_fixture_header(path: Path, cases: List[Dict[str, object]], metadata: Dict[str, object]) -> None: lines = [ "#pragma once", @@ -533,7 +540,7 @@ def run_fixtures(out_dir: Path) -> None: f"Missing checkpoint {checkpoint_path}. Run --mode train before --mode fixtures." ) - checkpoint = torch.load(checkpoint_path, map_location="cpu") + checkpoint = load_checkpoint(checkpoint_path) metadata = checkpoint["metadata"] print(f"[info] Loading checkpoint: {checkpoint_path}") print(f"[info] metadata={metadata}") From 839e2682930d9965dee8e87e1c482ea3356e06f0 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Mon, 1 Jun 2026 14:50:47 -0700 Subject: [PATCH 05/10] MAke diffusion MGN test not appear during build phase. --- tests/AMSlib/ams_interface/CMakeLists.txt | 19 +- .../test_graph_mgn_surrogate.cpp | 293 ++++++++++++++++-- tests/AMSlib/models/CMakeLists.txt | 89 +++++- .../models/generate_mgn_graph_diffusion.py | 227 ++++++-------- 4 files changed, 449 insertions(+), 179 deletions(-) diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index d8d2669f..d0f2b53b 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -34,7 +34,18 @@ BUILD_UNIT_TEST(ams_graph_surrogate test_graph_surrogate.cpp) ADD_AMS_UNIT_TEST(AMS_GRAPH_SURROGATE ams_graph_surrogate) BUILD_UNIT_TEST(ams_graph_mgn_surrogate test_graph_mgn_surrogate.cpp) -if (TARGET generate_mgn_graph_diffusion) - add_dependencies(ams_graph_mgn_surrogate generate_mgn_graph_diffusion) -endif() -ADD_AMS_UNIT_TEST(AMS_GRAPH_MGN_SURROGATE ams_graph_mgn_surrogate) +target_link_libraries(ams_graph_mgn_surrogate PRIVATE nlohmann_json::nlohmann_json) +target_compile_definitions(ams_graph_mgn_surrogate + PRIVATE + AMS_MGN_DIFFUSION_FIXTURE_DIR="${AMS_TEST_ROOT}/models/mgn_graph_diffusion" +) + +add_test( + NAME MGN_DIFFUSION_AMS_PARITY + COMMAND $ -s --reporter console +) +set_tests_properties(MGN_DIFFUSION_AMS_PARITY PROPERTIES + DEPENDS MGN_DIFFUSION_FIXTURES + LABELS "AMS_INTERFACE;MGN_DIFFUSION" + WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" +) diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp index 05f9a4c9..06466ab7 100644 --- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -1,21 +1,45 @@ #include #include +#include #include #include #include +#include +#include +#include +#include #include #include "AMS.h" #include "AMSGraph.hpp" #include "AMSTensor.hpp" -#include "models/mgn_graph_fixtures.hpp" using namespace ams; +using json = nlohmann::json; using Dim = AMSTensor::IntDimType; -static const char* MGN_GRAPH_MODEL_PATH = "../models/mgn_graph_diffusion.pt"; +#ifndef AMS_MGN_DIFFUSION_FIXTURE_DIR +#define AMS_MGN_DIFFUSION_FIXTURE_DIR "" +#endif + +namespace +{ + +constexpr int kFixtureFormatVersion = 1; +constexpr std::int64_t kNodeFeatureDim = 4; +constexpr std::int64_t kEdgeFeatureDim = 4; +constexpr std::int64_t kGlobalFeatureDim = 1; +constexpr std::int64_t kReferenceOutputDim = 1; + +struct TensorMetadata { + std::filesystem::path path; + std::string dtype; + std::vector shape; + std::string endianness; + std::uintmax_t byte_size; +}; static std::vector contiguousStrides(const std::vector& shape) { @@ -29,77 +53,286 @@ static std::vector contiguousStrides(const std::vector& shape) } template -static AMSTensor makeTensor(std::vector shape, const T* values) +static AMSTensor makeTensor(std::vector shape, const std::vector& values) { std::vector strides = contiguousStrides(shape); auto tensor = AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); - std::copy(values, values + tensor.elements(), tensor.template data()); + CATCH_REQUIRE(values.size() == static_cast(tensor.elements())); + std::copy(values.begin(), values.end(), tensor.template data()); return tensor; } -static Dim toDim(std::int64_t value) { return static_cast(value); } +static Dim toDim(std::int64_t value) +{ + CATCH_REQUIRE(value >= 0); + CATCH_REQUIRE(value <= std::numeric_limits::max()); + return static_cast(value); +} + +static bool hostIsLittleEndian() +{ + const std::uint16_t value = 1; + return *reinterpret_cast(&value) == 1; +} + +static std::uintmax_t dtypeByteWidth(const std::string& dtype) +{ + if (dtype == "float32") { + return sizeof(float); + } + if (dtype == "int64") { + return sizeof(std::int64_t); + } + CATCH_FAIL("Unsupported MGN graph diffusion tensor dtype in manifest: " + << dtype); + return 0; +} + +static std::uintmax_t shapeElementCount(const std::vector& shape) +{ + std::uintmax_t count = 1; + for (std::int64_t dim : shape) { + CATCH_REQUIRE(dim >= 0); + count *= static_cast(dim); + } + return count; +} + +static TensorMetadata parseTensorMetadata(const json& tensor) +{ + CATCH_REQUIRE(tensor.contains("path")); + CATCH_REQUIRE(tensor.contains("dtype")); + CATCH_REQUIRE(tensor.contains("shape")); + CATCH_REQUIRE(tensor.contains("endianness")); + CATCH_REQUIRE(tensor.contains("byte_size")); + + TensorMetadata metadata; + metadata.path = tensor.at("path").get(); + metadata.dtype = tensor.at("dtype").get(); + metadata.shape = tensor.at("shape").get>(); + metadata.endianness = tensor.at("endianness").get(); + metadata.byte_size = tensor.at("byte_size").get(); + return metadata; +} + +static void validateTensorMetadata(const TensorMetadata& metadata, + const std::string& expected_dtype, + const std::vector& expected_shape) +{ + CATCH_REQUIRE_FALSE(metadata.path.empty()); + CATCH_REQUIRE_FALSE(metadata.path.is_absolute()); + for (const auto& part : metadata.path) { + CATCH_REQUIRE(part.string() != ".."); + } + CATCH_REQUIRE(metadata.dtype == expected_dtype); + CATCH_REQUIRE(metadata.shape == expected_shape); + CATCH_REQUIRE(metadata.endianness == "little"); + CATCH_REQUIRE(hostIsLittleEndian()); -static AMSHomogeneousGraph makeGraph(const ams::test::mgn::GraphFixture& f) + const std::uintmax_t expected_size = + dtypeByteWidth(metadata.dtype) * shapeElementCount(metadata.shape); + CATCH_REQUIRE(metadata.byte_size == expected_size); +} + +template +static std::vector readTensorBinary(const std::filesystem::path& manifest_dir, + const json& tensor, + const std::string& expected_dtype, + const std::vector& expected_shape) { + TensorMetadata metadata = parseTensorMetadata(tensor); + validateTensorMetadata(metadata, expected_dtype, expected_shape); + + const std::filesystem::path tensor_path = manifest_dir / metadata.path; + CATCH_REQUIRE(std::filesystem::exists(tensor_path)); + CATCH_REQUIRE(std::filesystem::is_regular_file(tensor_path)); + CATCH_REQUIRE(std::filesystem::file_size(tensor_path) == metadata.byte_size); + + const std::uintmax_t element_count = shapeElementCount(metadata.shape); + CATCH_REQUIRE(element_count <= std::numeric_limits::max()); + std::vector values(static_cast(element_count)); + + std::ifstream input(tensor_path, std::ios::binary); + CATCH_REQUIRE(input); + if (metadata.byte_size > 0) { + CATCH_REQUIRE(metadata.byte_size <= + static_cast( + std::numeric_limits::max())); + input.read(reinterpret_cast(values.data()), + static_cast(metadata.byte_size)); + CATCH_REQUIRE(static_cast(input.gcount()) == + metadata.byte_size); + } + return values; +} + +static json loadManifest(const std::filesystem::path& fixture_dir) +{ + const std::filesystem::path manifest_path = fixture_dir / "fixtures.json"; + if (!std::filesystem::exists(manifest_path)) { + CATCH_FAIL("Missing MGN graph diffusion fixtures at " + << manifest_path + << ". Run `ctest -R MGN_DIFFUSION_FIXTURES` or the full " + "`ctest -R MGN_DIFFUSION` chain."); + } + + std::ifstream input(manifest_path); + CATCH_REQUIRE(input); + json manifest = json::parse(input); + CATCH_REQUIRE(manifest.contains("format_version")); + CATCH_REQUIRE(manifest.contains("endianness")); + CATCH_REQUIRE(manifest.at("format_version").get() == + kFixtureFormatVersion); + CATCH_REQUIRE(manifest.at("endianness").get() == "little"); + CATCH_REQUIRE(manifest.contains("model")); + CATCH_REQUIRE(manifest.contains("cases")); + CATCH_REQUIRE(manifest.contains("comparison")); + CATCH_REQUIRE(manifest.at("cases").is_array()); + return manifest; +} + +static std::filesystem::path resolveManifestRelativePath( + const std::filesystem::path& manifest_dir, const json& object) +{ + const std::filesystem::path relative_path = object.at("path").get(); + CATCH_REQUIRE_FALSE(relative_path.empty()); + CATCH_REQUIRE_FALSE(relative_path.is_absolute()); + for (const auto& part : relative_path) { + CATCH_REQUIRE(part.string() != ".."); + } + return manifest_dir / relative_path; +} + +static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, + const json& graph_case) +{ + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t num_edges = graph_case.at("num_edges").get(); + const std::int64_t node_dim = + graph_case.at("node_feature_dim").get(); + const std::int64_t edge_dim = + graph_case.at("edge_feature_dim").get(); + const std::int64_t global_dim = + graph_case.at("global_feature_dim").get(); + CATCH_REQUIRE(node_dim == kNodeFeatureDim); + CATCH_REQUIRE(edge_dim == kEdgeFeatureDim); + CATCH_REQUIRE(global_dim == kGlobalFeatureDim); + + const json& tensors = graph_case.at("tensors"); + auto node_features = readTensorBinary( + manifest_dir, tensors.at("node_features"), "float32", {num_nodes, node_dim}); + auto edge_index = readTensorBinary( + manifest_dir, tensors.at("edge_index"), "int64", {2, num_edges}); + auto edge_features = readTensorBinary( + manifest_dir, tensors.at("edge_features"), "float32", {num_edges, edge_dim}); + auto global_features = readTensorBinary( + manifest_dir, tensors.at("global_features"), "float32", {1, global_dim}); + return AMSHomogeneousGraph( - makeTensor({toDim(f.num_nodes), toDim(f.node_feature_dim)}, - f.node_features), - makeTensor({2, toDim(f.num_edges)}, f.edge_index), - makeTensor({toDim(f.num_edges), toDim(f.edge_feature_dim)}, - f.edge_features), - makeTensor({1, toDim(f.global_feature_dim)}, f.global_features)); + makeTensor({toDim(num_nodes), toDim(node_dim)}, node_features), + makeTensor({2, toDim(num_edges)}, edge_index), + makeTensor({toDim(num_edges), toDim(edge_dim)}, edge_features), + makeTensor({1, toDim(global_dim)}, global_features)); +} + +static std::vector loadReferenceDeltaU( + const std::filesystem::path& manifest_dir, const json& graph_case) +{ + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); + CATCH_REQUIRE(output_dim == kReferenceOutputDim); + return readTensorBinary(manifest_dir, + graph_case.at("tensors").at("reference_delta_u"), + "float32", + {num_nodes, output_dim}); } -static void verifyDeltaU(const ams::test::mgn::GraphFixture& f, - const AMSHomogeneousGraphFields& outputs) +static void verifyDeltaU(const json& graph_case, + const std::vector& reference_delta_u, + const AMSHomogeneousGraphFields& outputs, + double rtol, + double atol) { + const std::int64_t num_nodes = graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); + CATCH_REQUIRE(output_dim == kReferenceOutputDim); + CATCH_REQUIRE(outputs.node_fields.contains("delta_u")); const auto& delta_u = outputs.node_fields.at("delta_u"); - CATCH_REQUIRE(delta_u.shape()[0] == f.num_nodes); - CATCH_REQUIRE(delta_u.shape()[1] == f.reference_output_dim); + CATCH_REQUIRE(delta_u.shape()[0] == num_nodes); + CATCH_REQUIRE(delta_u.shape()[1] == output_dim); + CATCH_REQUIRE(reference_delta_u.size() == + static_cast(num_nodes * output_dim)); const float* actual = delta_u.data(); - const std::int64_t count = f.num_nodes * f.reference_output_dim; + const std::int64_t count = num_nodes * output_dim; for (std::int64_t i = 0; i < count; ++i) { - CATCH_REQUIRE(actual[i] == Catch::Approx(f.reference_delta_u[i]) - .epsilon(ams::test::mgn::kComparisonRtol) - .margin(ams::test::mgn::kComparisonAtol)); + CATCH_REQUIRE(actual[i] == Catch::Approx(reference_delta_u[i]) + .epsilon(rtol) + .margin(atol)); } } +} // namespace + CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", "[wf][graph][surrogate][mgn]") { - AMSInit(); + const std::filesystem::path fixture_dir = AMS_MGN_DIFFUSION_FIXTURE_DIR; + CATCH_REQUIRE_FALSE(fixture_dir.empty()); + const json manifest = loadManifest(fixture_dir); + const std::filesystem::path manifest_dir = fixture_dir; + const std::filesystem::path model_path = + resolveManifestRelativePath(manifest_dir, manifest.at("model")); + CATCH_REQUIRE(std::filesystem::exists(model_path)); + CATCH_REQUIRE(std::filesystem::is_regular_file(model_path)); + + const double rtol = manifest.at("comparison").at("rtol").get(); + const double atol = manifest.at("comparison").at("atol").get(); + CATCH_REQUIRE(manifest.at("cases").size() == 2); + + AMSInit(); + const std::string model_path_string = model_path.string(); auto model = AMSRegisterAbstractModel("test_mgn_graph_diffusion_surrogate", 0.5, - MGN_GRAPH_MODEL_PATH, + model_path_string.c_str(), false); AMSExecutor executor = AMSCreateExecutor(model, 0, 1); - for (std::size_t i = 0; i < ams::test::mgn::kNumGraphFixtures; ++i) { - const auto& fixture = ams::test::mgn::kGraphFixtures[i]; - CATCH_DYNAMIC_SECTION("fixture " << fixture.name) + const std::vector expected_node_counts = {24, 73}; + std::size_t case_index = 0; + for (const json& graph_case : manifest.at("cases")) { + CATCH_REQUIRE(graph_case.at("num_nodes").get() == + expected_node_counts.at(case_index)); + CATCH_DYNAMIC_SECTION("fixture " << graph_case.at("name").get()) { - AMSHomogeneousGraph graph = makeGraph(fixture); + AMSHomogeneousGraph graph = makeGraph(manifest_dir, graph_case); + std::vector reference_delta_u = + loadReferenceDeltaU(manifest_dir, graph_case); bool callback_invoked = false; HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&, AMSHomogeneousGraphFields& outputs) { callback_invoked = true; + const std::int64_t num_nodes = + graph_case.at("num_nodes").get(); + const std::int64_t output_dim = + graph_case.at("reference_output_dim").get(); outputs.node_fields.set( "delta_u", - makeTensor({toDim(fixture.num_nodes), - toDim(fixture.reference_output_dim)}, - fixture.reference_delta_u)); + makeTensor({toDim(num_nodes), toDim(output_dim)}, + reference_delta_u)); }; AMSHomogeneousGraphFields outputs; AMSExecute(executor, callback, graph, outputs); CATCH_REQUIRE_FALSE(callback_invoked); - verifyDeltaU(fixture, outputs); + verifyDeltaU(graph_case, reference_delta_u, outputs, rtol, atol); } + ++case_index; } } diff --git a/tests/AMSlib/models/CMakeLists.txt b/tests/AMSlib/models/CMakeLists.txt index e3f874f3..a9cbf950 100644 --- a/tests/AMSlib/models/CMakeLists.txt +++ b/tests/AMSlib/models/CMakeLists.txt @@ -118,41 +118,66 @@ configure_file( find_program(AMS_PYTHON_EXECUTABLE NAMES python3 python REQUIRED) +set(MGN_GRAPH_DIFFUSION_DIR + "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion" +) set(MGN_GRAPH_DIFFUSION_CHECKPOINT - "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion_checkpoint.pt" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion_checkpoint.pt" ) set(MGN_GRAPH_DIFFUSION_MODEL - "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion.pt" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion.pt" +) +set(MGN_GRAPH_DIFFUSION_FEASIBILITY_MODEL + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_graph_diffusion_feasibility.pt" +) +set(MGN_GRAPH_DIFFUSION_MANIFEST + "${MGN_GRAPH_DIFFUSION_DIR}/fixtures.json" +) +set(MGN_GRAPH_DIFFUSION_METRICS + "${MGN_GRAPH_DIFFUSION_DIR}/training_metrics.json" ) -set(MGN_GRAPH_DIFFUSION_FIXTURES - "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_fixtures.hpp" +set(MGN_GRAPH_DIFFUSION_TENSOR_FILES + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/node_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/edge_index.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/edge_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/global_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n24/reference_delta_u.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/node_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/edge_index.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/edge_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/global_features.bin" + "${MGN_GRAPH_DIFFUSION_DIR}/mgn_diffusion_n73/reference_delta_u.bin" ) add_custom_command( OUTPUT "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" "${MGN_GRAPH_DIFFUSION_MODEL}" - "${MGN_GRAPH_DIFFUSION_FIXTURES}" + "${MGN_GRAPH_DIFFUSION_MANIFEST}" + ${MGN_GRAPH_DIFFUSION_TENSOR_FILES} + BYPRODUCTS + "${MGN_GRAPH_DIFFUSION_FEASIBILITY_MODEL}" + "${MGN_GRAPH_DIFFUSION_METRICS}" COMMAND "${AMS_PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" - --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" --mode feasibility COMMAND "${AMS_PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" - --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" --mode train COMMAND "${AMS_PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" - --out-dir "${CMAKE_CURRENT_BINARY_DIR}" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" --mode fixtures DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" "${CMAKE_CURRENT_SOURCE_DIR}/ams_model.py" - WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}" - COMMENT "Generating Stage 2 pure-Torch MGN graph diffusion model and fixtures" + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" + COMMENT "Generating pure-Torch MGN graph diffusion model and fixtures" VERBATIM ) @@ -160,5 +185,47 @@ add_custom_target(generate_mgn_graph_diffusion DEPENDS "${MGN_GRAPH_DIFFUSION_CHECKPOINT}" "${MGN_GRAPH_DIFFUSION_MODEL}" - "${MGN_GRAPH_DIFFUSION_FIXTURES}" + "${MGN_GRAPH_DIFFUSION_MANIFEST}" + ${MGN_GRAPH_DIFFUSION_TENSOR_FILES} +) + +add_test( + NAME MGN_DIFFUSION_FEASIBILITY + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode feasibility +) +set_tests_properties(MGN_DIFFUSION_FEASIBILITY PROPERTIES + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" +) + +add_test( + NAME MGN_DIFFUSION_TRAIN + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode train +) +set_tests_properties(MGN_DIFFUSION_TRAIN PROPERTIES + DEPENDS MGN_DIFFUSION_FEASIBILITY + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" +) + +add_test( + NAME MGN_DIFFUSION_FIXTURES + COMMAND + "${AMS_PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/generate_mgn_graph_diffusion.py" + --out-dir "${MGN_GRAPH_DIFFUSION_DIR}" + --mode fixtures +) +set_tests_properties(MGN_DIFFUSION_FIXTURES PROPERTIES + DEPENDS MGN_DIFFUSION_TRAIN + LABELS MGN_DIFFUSION + WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" ) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 5c05a5bc..1e868bd5 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -1,15 +1,17 @@ #!/usr/bin/env python3 """Train/export a tiny pure-Torch MGN-like graph diffusion surrogate. -This Stage 2 utility intentionally stays pure PyTorch. It generates synthetic -homogeneous graphs, trains a small message-passing model, exports an -AMS-wrapped TorchScript model, and writes C++ fixtures whose reference outputs -come from the exported TorchScript model. +This MGN graph diffusion utility intentionally stays pure PyTorch. It generates +synthetic homogeneous graphs, trains a small message-passing model, exports an +AMS-wrapped TorchScript model, and writes runtime fixtures whose reference +outputs come from the exported TorchScript model. """ import argparse import copy +import json import os +import struct import sys from pathlib import Path from typing import Dict, Iterable, List, Tuple @@ -20,7 +22,7 @@ from torch import Tensor except ModuleNotFoundError as exc: raise SystemExit( - "PyTorch is required for Stage 2 MGN graph diffusion generation. " + "PyTorch is required for MGN graph diffusion generation. " "Install PyTorch or run this script on a Torch-enabled machine." ) from exc @@ -55,7 +57,12 @@ CHECKPOINT_NAME = "mgn_graph_diffusion_checkpoint.pt" MODEL_NAME = "mgn_graph_diffusion.pt" -FIXTURE_HEADER_NAME = "mgn_graph_fixtures.hpp" +FIXTURE_MANIFEST_NAME = "fixtures.json" +TRAINING_METRICS_NAME = "training_metrics.json" +FIXTURE_FORMAT_VERSION = 1 +FIXTURE_ENDIANNESS = "little" +COMPARISON_RTOL = 2.0e-5 +COMPARISON_ATOL = 2.0e-5 class MLP(nn.Module): @@ -248,7 +255,7 @@ def assert_close(name: str, actual: Tensor, expected: Tensor, atol: float = 1.0e def run_feasibility(out_dir: Path) -> None: out_dir.mkdir(parents=True, exist_ok=True) - print("[info] Stage 2 feasibility gate") + print("[info] MGN graph diffusion feasibility gate") print(f"[info] model_seed={MODEL_SEED}, fixture_sizes={FIXTURE_GRAPH_SIZES}, config={model_config()}") model = make_model().eval() @@ -308,7 +315,7 @@ def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]] def run_train(out_dir: Path) -> None: out_dir.mkdir(parents=True, exist_ok=True) - print("[info] Stage 2 training") + print("[info] MGN graph diffusion training") print(f"[info] model_seed={MODEL_SEED}, train_data_seed={TRAIN_DATA_SEED}, val_data_seed={VAL_DATA_SEED}") print(f"[info] config={model_config()}") print( @@ -376,7 +383,7 @@ def run_train(out_dir: Path) -> None: ) if val_mse > primary_threshold: raise RuntimeError( - "Training failed primary Stage 2 criterion: " + "Training failed primary MGN graph diffusion criterion: " f"validation_mse={val_mse:.8e} > {primary_threshold:.8e}" ) @@ -410,33 +417,9 @@ def run_train(out_dir: Path) -> None: torch.save(checkpoint, checkpoint_path) print(f"[info] Saved checkpoint: {checkpoint_path}") - -def format_float(value: float) -> str: - text = f"{float(value):.9g}" - if "e" not in text and "." not in text: - text += ".0" - return text + "f" - - -def format_array(name: str, c_type: str, values: List[float], values_per_line: int = 6) -> str: - lines = [f"inline constexpr {c_type} {name}[] = {{"] - for i in range(0, len(values), values_per_line): - chunk = values[i : i + values_per_line] - if c_type == "float": - rendered = ", ".join(format_float(v) for v in chunk) - else: - rendered = ", ".join(str(int(v)) for v in chunk) - lines.append(f" {rendered},") - lines.append("};") - return "\n".join(lines) - - -def flatten_float(tensor: Tensor) -> List[float]: - return [float(x) for x in tensor.detach().cpu().contiguous().view(-1).tolist()] - - -def flatten_int(tensor: Tensor) -> List[int]: - return [int(x) for x in tensor.detach().cpu().contiguous().view(-1).tolist()] + metrics_path = out_dir / TRAINING_METRICS_NAME + metrics_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"[info] Saved training metrics: {metrics_path}") def load_checkpoint(path: Path) -> Dict[str, object]: @@ -446,90 +429,32 @@ def load_checkpoint(path: Path) -> Dict[str, object]: return torch.load(path, map_location="cpu") -def write_fixture_header(path: Path, cases: List[Dict[str, object]], metadata: Dict[str, object]) -> None: - lines = [ - "#pragma once", - "", - "#include ", - "#include ", - "", - "// Generated by tests/AMSlib/models/generate_mgn_graph_diffusion.py --mode fixtures.", - f"// model_seed: {metadata['model_seed']}", - f"// train_data_seed: {metadata['train_data_seed']}", - f"// val_data_seed: {metadata['val_data_seed']}", - f"// fixture_graph_sizes: {metadata['fixture_graph_sizes']}", - f"// fixture_seeds: {metadata['fixture_seeds']}", - f"// model_config: {metadata['model_config']}", - f"// validation_mse: {metadata['validation_mse']:.9e}", - f"// zero_baseline_mse: {metadata['zero_baseline_mse']:.9e}", - f"// primary criterion: validation_mse <= {metadata['primary_baseline_factor']} * zero_baseline_mse", - "", - "namespace ams::test::mgn", - "{", - "", - "struct GraphFixture {", - " const char* name;", - " std::int64_t num_nodes;", - " std::int64_t num_edges;", - " std::int64_t node_feature_dim;", - " std::int64_t edge_feature_dim;", - " std::int64_t global_feature_dim;", - " std::int64_t reference_output_dim;", - " const float* node_features;", - " const std::int64_t* edge_index;", - " const float* edge_features;", - " const float* global_features;", - " const float* reference_delta_u;", - "};", - "", - "inline constexpr float kComparisonRtol = 2.0e-5f;", - "inline constexpr float kComparisonAtol = 2.0e-5f;", - "", - ] - - for case in cases: - prefix = str(case["prefix"]) - lines.append(f"// {case['name']}: N={case['num_nodes']}, E={case['num_edges']}") - lines.append(format_array(f"{prefix}_node_features", "float", case["node_features"])) - lines.append(format_array(f"{prefix}_edge_index", "std::int64_t", case["edge_index"], values_per_line=8)) - lines.append(format_array(f"{prefix}_edge_features", "float", case["edge_features"])) - lines.append(format_array(f"{prefix}_global_features", "float", case["global_features"])) - lines.append(format_array(f"{prefix}_reference_delta_u", "float", case["reference_delta_u"])) - lines.append("") - - lines.append("inline constexpr GraphFixture kGraphFixtures[] = {") - for case in cases: - prefix = str(case["prefix"]) - lines.extend( - [ - " {", - f" \"{case['name']}\",", - f" {case['num_nodes']},", - f" {case['num_edges']},", - f" {NODE_FEATURE_DIM},", - f" {EDGE_FEATURE_DIM},", - f" {GLOBAL_FEATURE_DIM},", - f" {REFERENCE_OUTPUT_DIM},", - f" {prefix}_node_features,", - f" {prefix}_edge_index,", - f" {prefix}_edge_features,", - f" {prefix}_global_features,", - f" {prefix}_reference_delta_u,", - " },", - ] - ) - lines.extend( - [ - "};", - "", - "inline constexpr std::size_t kNumGraphFixtures =", - " sizeof(kGraphFixtures) / sizeof(kGraphFixtures[0]);", - "", - "} // namespace ams::test::mgn", - "", - ] - ) - path.write_text("\n".join(lines), encoding="utf-8") +def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype: str) -> Dict[str, object]: + rel_path = Path(relative_path) + if rel_path.is_absolute() or ".." in rel_path.parts: + raise ValueError(f"Fixture tensor path must stay relative: {relative_path}") + + if dtype == "float32": + packed_tensor = tensor.detach().cpu().to(dtype=torch.float32).contiguous() + values = [float(x) for x in packed_tensor.view(-1).tolist()] + data = struct.pack("<" + "f" * len(values), *values) if values else b"" + elif dtype == "int64": + packed_tensor = tensor.detach().cpu().to(dtype=torch.int64).contiguous() + values = [int(x) for x in packed_tensor.view(-1).tolist()] + data = struct.pack("<" + "q" * len(values), *values) if values else b"" + else: + raise ValueError(f"Unsupported fixture tensor dtype: {dtype}") + + tensor_path = out_dir / rel_path + tensor_path.parent.mkdir(parents=True, exist_ok=True) + tensor_path.write_bytes(data) + return { + "path": rel_path.as_posix(), + "dtype": dtype, + "shape": [int(dim) for dim in packed_tensor.shape], + "endianness": FIXTURE_ENDIANNESS, + "byte_size": len(data), + } def run_fixtures(out_dir: Path) -> None: @@ -537,7 +462,8 @@ def run_fixtures(out_dir: Path) -> None: checkpoint_path = out_dir / CHECKPOINT_NAME if not checkpoint_path.exists(): raise RuntimeError( - f"Missing checkpoint {checkpoint_path}. Run --mode train before --mode fixtures." + f"Missing checkpoint {checkpoint_path}. Run --mode train, " + "or `ctest -R MGN_DIFFUSION_TRAIN`, before --mode fixtures." ) checkpoint = load_checkpoint(checkpoint_path) @@ -556,20 +482,39 @@ def run_fixtures(out_dir: Path) -> None: reloaded = torch.jit.load(str(model_path)).eval() cases: List[Dict[str, object]] = [] with torch.no_grad(): - for index, (num_nodes, seed) in enumerate(zip(FIXTURE_GRAPH_SIZES, FIXTURE_SEEDS)): + for num_nodes, seed in zip(FIXTURE_GRAPH_SIZES, FIXTURE_SEEDS): graph, _target = generate_graph(num_nodes, seed) graph = graph_to_model_input(graph) output = reloaded(graph)["node:delta_u"].detach().cpu().contiguous() + name = f"mgn_diffusion_n{num_nodes}" + case_dir = name + tensors = { + "node_features": write_tensor_binary( + out_dir, f"{case_dir}/node_features.bin", graph["node_features"], "float32" + ), + "edge_index": write_tensor_binary( + out_dir, f"{case_dir}/edge_index.bin", graph["edge_index"], "int64" + ), + "edge_features": write_tensor_binary( + out_dir, f"{case_dir}/edge_features.bin", graph["edge_features"], "float32" + ), + "global_features": write_tensor_binary( + out_dir, f"{case_dir}/global_features.bin", graph["global_features"], "float32" + ), + "reference_delta_u": write_tensor_binary( + out_dir, f"{case_dir}/reference_delta_u.bin", output, "float32" + ), + } case = { - "name": f"mgn_diffusion_n{num_nodes}", - "prefix": f"kCase{index}", + "name": name, + "seed": int(seed), "num_nodes": int(num_nodes), "num_edges": int(graph["edge_index"].shape[1]), - "node_features": flatten_float(graph["node_features"]), - "edge_index": flatten_int(graph["edge_index"]), - "edge_features": flatten_float(graph["edge_features"]), - "global_features": flatten_float(graph["global_features"]), - "reference_delta_u": flatten_float(output), + "node_feature_dim": NODE_FEATURE_DIM, + "edge_feature_dim": EDGE_FEATURE_DIM, + "global_feature_dim": GLOBAL_FEATURE_DIM, + "reference_output_dim": REFERENCE_OUTPUT_DIM, + "tensors": tensors, } print( f"[info] fixture {case['name']} N={case['num_nodes']} " @@ -577,19 +522,33 @@ def run_fixtures(out_dir: Path) -> None: ) cases.append(case) - header_path = out_dir / FIXTURE_HEADER_NAME - write_fixture_header(header_path, cases, metadata) - print(f"[info] Saved C++ fixture header: {header_path}") + manifest = { + "format_version": FIXTURE_FORMAT_VERSION, + "endianness": FIXTURE_ENDIANNESS, + "model": { + "path": MODEL_NAME, + }, + "training_metrics_path": TRAINING_METRICS_NAME, + "metadata": metadata, + "comparison": { + "rtol": COMPARISON_RTOL, + "atol": COMPARISON_ATOL, + }, + "cases": cases, + } + manifest_path = out_dir / FIXTURE_MANIFEST_NAME + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"[info] Saved fixture manifest: {manifest_path}") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--out-dir", type=Path, required=True, help="Output directory for Stage 2 artifacts") + parser.add_argument("--out-dir", type=Path, required=True, help="Output directory for MGN graph diffusion artifacts") parser.add_argument( "--mode", choices=("feasibility", "train", "fixtures", "all"), required=True, - help="Separated Stage 2 generation mode", + help="Separated MGN graph diffusion generation mode", ) args = parser.parse_args() From 278167ccd885352ed7c9f00dac7a4439fc048ae7 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Mon, 1 Jun 2026 14:58:59 -0700 Subject: [PATCH 06/10] Add documentation. --- tests/AMSlib/ams_interface/CMakeLists.txt | 2 + .../test_graph_mgn_surrogate.cpp | 26 +++++ tests/AMSlib/models/CMakeLists.txt | 5 + .../models/generate_mgn_graph_diffusion.py | 96 ++++++++++++++++++- 4 files changed, 125 insertions(+), 4 deletions(-) diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index d0f2b53b..38d89399 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -40,6 +40,8 @@ target_compile_definitions(ams_graph_mgn_surrogate AMS_MGN_DIFFUSION_FIXTURE_DIR="${AMS_TEST_ROOT}/models/mgn_graph_diffusion" ) +# Final runtime check for the MGN diffusion workflow: fixtures mode must have +# produced the TorchScript model and binary graph cases before AMS parity runs. add_test( NAME MGN_DIFFUSION_AMS_PARITY COMMAND $ -s --reporter console diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp index 06466ab7..184ee7ea 100644 --- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -20,6 +20,10 @@ using json = nlohmann::json; using Dim = AMSTensor::IntDimType; +// The MGN diffusion workflow writes fixtures into the build tree at CTest +// runtime. CMake compiles that build-tree directory into this focused parity +// test so the test can be launched like any other Catch executable while still +// keeping generated tensors out of the source tree. #ifndef AMS_MGN_DIFFUSION_FIXTURE_DIR #define AMS_MGN_DIFFUSION_FIXTURE_DIR "" #endif @@ -119,6 +123,9 @@ static void validateTensorMetadata(const TensorMetadata& metadata, const std::string& expected_dtype, const std::vector& expected_shape) { + // Validate the manifest before allocating AMSTensors. The binary files are + // intentionally simple raw bytes, so the manifest is the only place where + // dtype, shape, byte order, and size are described. CATCH_REQUIRE_FALSE(metadata.path.empty()); CATCH_REQUIRE_FALSE(metadata.path.is_absolute()); for (const auto& part : metadata.path) { @@ -143,6 +150,8 @@ static std::vector readTensorBinary(const std::filesystem::path& manifest_dir TensorMetadata metadata = parseTensorMetadata(tensor); validateTensorMetadata(metadata, expected_dtype, expected_shape); + // All tensor paths are relative to fixtures.json. That keeps the generated + // fixture directory relocatable inside different build trees. const std::filesystem::path tensor_path = manifest_dir / metadata.path; CATCH_REQUIRE(std::filesystem::exists(tensor_path)); CATCH_REQUIRE(std::filesystem::is_regular_file(tensor_path)); @@ -179,6 +188,9 @@ static json loadManifest(const std::filesystem::path& fixture_dir) std::ifstream input(manifest_path); CATCH_REQUIRE(input); json manifest = json::parse(input); + // Keep the manifest version check close to parsing so future fixture format + // changes fail loudly instead of being interpreted as the current raw-binary + // contract. CATCH_REQUIRE(manifest.contains("format_version")); CATCH_REQUIRE(manifest.contains("endianness")); CATCH_REQUIRE(manifest.at("format_version").get() == @@ -194,6 +206,8 @@ static json loadManifest(const std::filesystem::path& fixture_dir) static std::filesystem::path resolveManifestRelativePath( const std::filesystem::path& manifest_dir, const json& object) { + // Model and tensor paths in the manifest are relative by design. Avoiding + // absolute paths makes fixtures reusable if the whole build directory moves. const std::filesystem::path relative_path = object.at("path").get(); CATCH_REQUIRE_FALSE(relative_path.empty()); CATCH_REQUIRE_FALSE(relative_path.is_absolute()); @@ -206,6 +220,9 @@ static std::filesystem::path resolveManifestRelativePath( static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, const json& graph_case) { + // Runtime fixture binaries become the same AMSTensor-backed homogeneous graph + // that an application would pass to AMS: node features, canonical edge_index, + // edge features, and required global features. const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t num_edges = graph_case.at("num_edges").get(); const std::int64_t node_dim = @@ -254,6 +271,8 @@ static void verifyDeltaU(const json& graph_case, double rtol, double atol) { + // The reference is the Python TorchScript output, not the exact synthetic + // diffusion target. This test is about AMS/LibTorch deployment parity. const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t output_dim = graph_case.at("reference_output_dim").get(); @@ -302,6 +321,9 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", false); AMSExecutor executor = AMSCreateExecutor(model, 0, 1); + // The two fixed fixture sizes intentionally have different N and E values. + // Running both cases catches accidental static-shape assumptions in the + // exported TorchScript model or in AMS graph tensor handling. const std::vector expected_node_counts = {24, 73}; std::size_t case_index = 0; for (const json& graph_case : manifest.at("cases")) { @@ -314,6 +336,10 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", loadReferenceDeltaU(manifest_dir, graph_case); bool callback_invoked = false; + // If the surrogate path fails, AMS would call the domain fallback. For + // this parity test that would hide a deployment failure, so the callback + // records whether it was invoked and supplies the reference only as a + // valid fallback value. HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&, AMSHomogeneousGraphFields& outputs) { callback_invoked = true; diff --git a/tests/AMSlib/models/CMakeLists.txt b/tests/AMSlib/models/CMakeLists.txt index a9cbf950..ec40888d 100644 --- a/tests/AMSlib/models/CMakeLists.txt +++ b/tests/AMSlib/models/CMakeLists.txt @@ -118,6 +118,9 @@ configure_file( find_program(AMS_PYTHON_EXECUTABLE NAMES python3 python REQUIRED) +# MGN diffusion artifacts are generated in the build tree. The custom target is +# available for explicit regeneration, but normal C++ builds do not depend on +# training/export; the CTest workflow below owns the runtime validation path. set(MGN_GRAPH_DIFFUSION_DIR "${CMAKE_CURRENT_BINARY_DIR}/mgn_graph_diffusion" ) @@ -189,6 +192,8 @@ add_custom_target(generate_mgn_graph_diffusion ${MGN_GRAPH_DIFFUSION_TENSOR_FILES} ) +# Runtime workflow tests. The dependency chain keeps the expensive learned-model +# path ordered even under parallel CTest: feasibility -> train -> fixtures. add_test( NAME MGN_DIFFUSION_FEASIBILITY COMMAND diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 1e868bd5..5ccdcfca 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -1,10 +1,32 @@ #!/usr/bin/env python3 """Train/export a tiny pure-Torch MGN-like graph diffusion surrogate. -This MGN graph diffusion utility intentionally stays pure PyTorch. It generates -synthetic homogeneous graphs, trains a small message-passing model, exports an -AMS-wrapped TorchScript model, and writes runtime fixtures whose reference -outputs come from the exported TorchScript model. +This utility is a self-contained learned-graph smoke test for the AMS graph +surrogate path. It intentionally stays pure PyTorch: no PyG/DGL/PhysicsNeMo, +no MFEM data, and no extra graph-learning runtime dependencies. + +The workflow is split into three explicit modes: + +* feasibility: create the model, run eager inference on two graph sizes, script + it, reload it, and verify all outputs match. This gates TorchScript and + dynamic N/E support before any training cost is paid. +* train: train the tiny model on deterministic synthetic graph diffusion data + and save only a checkpoint/state dict plus training metrics. +* fixtures: load the checkpoint, export the AMS-wrapped TorchScript model, and + write runtime fixture files consumed by the C++ parity test. + +The graph contract matches AMS homogeneous graph inputs: + + node_features [N, 4] = x, y, u, kappa + edge_index [2, E] = source row, destination row + edge_features [E, 4] = dx, dy, distance, normalized diffusion message + global_features [1, 1] = dt + +The exact target is a deterministic diffusion-like update, +delta_u_i = dt * sum_{src -> i} weight * (u_src - u_i). Training uses that +target, while the C++ fixture reference uses the exported TorchScript model +output. That distinction is deliberate: the AMS test validates deployment +parity with Python TorchScript, not physical fidelity of the learned model. """ import argparse @@ -79,6 +101,8 @@ def forward(self, x: Tensor) -> Tensor: class ProcessorBlock(nn.Module): + """One MGN-like processor step implemented with plain tensor operations.""" + def __init__(self, latent_dim: int, global_dim: int): super().__init__() self.edge_mlp = MLP(latent_dim * 3 + global_dim, latent_dim, latent_dim) @@ -91,9 +115,15 @@ def forward( edge_index: Tensor, global_features: Tensor, ) -> Tuple[Tensor, Tensor]: + # The canonical AMS edge_index uses row 0 as source and row 1 as + # destination. index_select gathers per-edge source/destination node + # states without any graph-library dependency. src = edge_index[0] dst = edge_index[1] + # Edge update: each edge sees its current latent state, source node, + # destination node, and the graph-level dt feature. The residual update + # keeps this tiny model easy to train. global_edges = global_features.expand(edge_latent.shape[0], global_features.shape[1]) edge_input = torch.cat( ( @@ -106,6 +136,9 @@ def forward( ) edge_latent = edge_latent + self.edge_mlp(edge_input) + # Node aggregation: incoming edge states are summed onto destination + # nodes with index_add, the same primitive TorchScript and LibTorch will + # execute after export. aggregated = torch.zeros( (node_latent.shape[0], edge_latent.shape[1]), dtype=node_latent.dtype, @@ -113,6 +146,8 @@ def forward( ) aggregated = aggregated.index_add(0, dst, edge_latent) + # Node update: each node sees its previous latent state, the aggregated + # incoming message, and the global dt feature. global_nodes = global_features.expand(node_latent.shape[0], global_features.shape[1]) node_input = torch.cat((node_latent, aggregated, global_nodes), dim=1) node_latent = node_latent + self.node_mlp(node_input) @@ -120,6 +155,8 @@ def forward( class TinyGraphDiffusionMGN(nn.Module): + """Small fixed-shape MeshGraphNet-like model for the AMS graph contract.""" + def __init__( self, node_dim: int = NODE_FEATURE_DIM, @@ -135,6 +172,9 @@ def __init__( self.node_decoder = MLP(latent_dim, latent_dim, REFERENCE_OUTPUT_DIM) def forward(self, graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + # Keep this signature and return type aligned with the AMS homogeneous + # graph surrogate contract. The output key becomes outputs.node_fields + # entry "delta_u" in C++. node_features = graph["node_features"] edge_index = graph["edge_index"].to(torch.int64) edge_features = graph["edge_features"] @@ -177,6 +217,9 @@ def make_generator(seed: int) -> torch.Generator: def _unique_directed_edges(knn: Tensor) -> Tensor: + # kNN is computed from each node's perspective. Adding both directions makes + # the fixture exercise directed source/destination semantics while still + # representing an undirected diffusion neighborhood. pairs = set() num_nodes = int(knn.shape[0]) for dst in range(num_nodes): @@ -191,8 +234,14 @@ def _unique_directed_edges(knn: Tensor) -> Tensor: def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor]: + """Create one deterministic synthetic graph and its exact diffusion target.""" + generator = make_generator(seed) positions = torch.rand((num_nodes, 2), generator=generator, dtype=torch.float32) + + # Use a smooth random field rather than independent random u values. This is + # still synthetic and cheap, but it looks more like a heat-equation state and + # keeps target scales stable enough to avoid feature/target normalization. phases = torch.rand((1, 4), generator=generator, dtype=torch.float32) * (2.0 * torch.pi) amps = torch.rand((1, 4), generator=generator, dtype=torch.float32) * 0.5 + 0.5 x = positions[:, 0:1] @@ -207,6 +256,9 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 dt = torch.rand((1, 1), generator=generator, dtype=torch.float32) * 0.06 + 0.02 + # Build a small directed kNN graph. E varies with N because duplicate edges + # are removed after adding the reverse direction, so N=24 and N=73 exercise + # dynamic node and edge counts in TorchScript and C++. distances = torch.cdist(positions, positions) masked = distances + torch.eye(num_nodes, dtype=torch.float32) * 1.0e6 knn = torch.topk(masked, k=K_NEIGHBORS, largest=False, dim=1).indices @@ -218,13 +270,24 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor distance = delta_pos.norm(dim=1, keepdim=True).clamp_min(1.0e-6) conductivity = 0.5 * (kappa.index_select(0, src) + kappa.index_select(0, dst)) raw_weight = conductivity / (distance + 0.05) + + # Normalize incoming weights per destination node. This is edge-weight + # normalization for a bounded synthetic target, not feature or target + # normalization embedded in the model. incoming_sum = torch.zeros((num_nodes, 1), dtype=torch.float32) incoming_sum = incoming_sum.index_add(0, dst, raw_weight) weight = raw_weight / incoming_sum.index_select(0, dst).clamp_min(1.0e-8) + # The fourth edge feature is the already weighted scalar message. Including + # it keeps the learned problem small while still requiring source/destination + # indexing and destination aggregation to recover node:delta_u. node_features = torch.cat((positions, u, kappa), dim=1).contiguous() messages = weight * (u.index_select(0, src) - u.index_select(0, dst)) edge_features = torch.cat((delta_pos, distance, messages), dim=1).contiguous() + + # Exact supervised target: aggregate incoming weighted messages and scale by + # dt. Python training checks the model against this target; C++ checks AMS + # against TorchScript reference outputs generated after export. target = torch.zeros((num_nodes, 1), dtype=torch.float32) target = target.index_add(0, dst, messages) target = dt * target @@ -239,6 +302,8 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor def graph_to_model_input(graph: Dict[str, Tensor]) -> Dict[str, Tensor]: + """Canonicalize dtypes before eager, scripted, or reloaded model calls.""" + return { "node_features": graph["node_features"].to(dtype=torch.float32), "edge_index": graph["edge_index"].to(dtype=torch.int64), @@ -254,6 +319,9 @@ def assert_close(name: str, actual: Tensor, expected: Tensor, atol: float = 1.0e def run_feasibility(out_dir: Path) -> None: + # This mode owns only TorchScript compatibility. It intentionally does not + # train or write deployment fixtures, so failures point at scripting, + # dynamic graph sizes, or AMS wrapper compatibility. out_dir.mkdir(parents=True, exist_ok=True) print("[info] MGN graph diffusion feasibility gate") print(f"[info] model_seed={MODEL_SEED}, fixture_sizes={FIXTURE_GRAPH_SIZES}, config={model_config()}") @@ -288,6 +356,8 @@ def run_feasibility(out_dir: Path) -> None: def validation_graphs() -> List[Tuple[Dict[str, Tensor], Tensor]]: + # Fixed validation seeds make the primary "10x better than zero baseline" + # criterion reproducible across regeneration runs. graphs = [] for i in range(VAL_GRAPHS): num_nodes = 16 + ((i * 37) % 113) @@ -296,6 +366,8 @@ def validation_graphs() -> List[Tuple[Dict[str, Tensor], Tensor]]: def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]]) -> Tuple[float, float, Tensor]: + """Return model MSE, zero-prediction baseline MSE, and all targets.""" + total_loss = 0.0 total_baseline = 0.0 total_nodes = 0 @@ -314,6 +386,8 @@ def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]] def run_train(out_dir: Path) -> None: + # This mode owns learned weights and metrics only. It writes a checkpoint + # that fixtures mode must consume; fixtures mode never silently retrains. out_dir.mkdir(parents=True, exist_ok=True) print("[info] MGN graph diffusion training") print(f"[info] model_seed={MODEL_SEED}, train_data_seed={TRAIN_DATA_SEED}, val_data_seed={VAL_DATA_SEED}") @@ -423,6 +497,9 @@ def run_train(out_dir: Path) -> None: def load_checkpoint(path: Path) -> Dict[str, object]: + # PyTorch 2.6 changed torch.load's default weights_only behavior. The + # checkpoint is generated by this local script and contains metadata, so + # loading with weights_only=False is intentional here. try: return torch.load(path, map_location="cpu", weights_only=False) except TypeError: @@ -430,6 +507,13 @@ def load_checkpoint(path: Path) -> Dict[str, object]: def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype: str) -> Dict[str, object]: + """Write one fixture tensor and return the matching manifest entry. + + The C++ test deliberately avoids NumPy or torch file readers. Each tensor is + raw contiguous row-major little-endian bytes, and fixtures.json records the + dtype, shape, endianness, and byte count needed to validate it before use. + """ + rel_path = Path(relative_path) if rel_path.is_absolute() or ".." in rel_path.parts: raise ValueError(f"Fixture tensor path must stay relative: {relative_path}") @@ -458,6 +542,10 @@ def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype def run_fixtures(out_dir: Path) -> None: + # This mode owns deployment artifacts: the AMS-wrapped TorchScript model, + # the fixture manifest, and raw tensor binaries. Reference outputs are model + # outputs from the reloaded TorchScript artifact, because C++ parity should + # answer "does AMS reproduce Python TorchScript inference?" out_dir.mkdir(parents=True, exist_ok=True) checkpoint_path = out_dir / CHECKPOINT_NAME if not checkpoint_path.exists(): From 59905974e41925eaee4b418a9486ca8136731c43 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Mon, 8 Jun 2026 11:54:49 -0700 Subject: [PATCH 07/10] clang-format. --- .../ams_interface/test_graph_fallback.cpp | 20 +---- .../test_graph_mgn_surrogate.cpp | 81 +++++++++++-------- 2 files changed, 52 insertions(+), 49 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_fallback.cpp b/tests/AMSlib/ams_interface/test_graph_fallback.cpp index 16e6e75f..b261e178 100644 --- a/tests/AMSlib/ams_interface/test_graph_fallback.cpp +++ b/tests/AMSlib/ams_interface/test_graph_fallback.cpp @@ -124,21 +124,9 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", AMSInit(); CATCH_REQUIRE_NOTHROW(makeValidGraph()); - { - AMSHomogeneousGraph graph(makeNodeFeatures(), - makeEdgeIndex32(), - makeEdgeFeatures()); - CATCH_REQUIRE(graph.global_features.shape().size() == 1); - CATCH_REQUIRE(graph.global_features.shape()[0] == 0); - } - CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeTensor({0}))); CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeGlobalFeatures())); + makeEdgeIndex32(), + makeEdgeFeatures())); CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeTensor({3, 2}), makeEdgeIndex64(), @@ -175,7 +163,7 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), makeEdgeFeatures(), - makeTensor({1, 2})), + makeTensor({2})), std::runtime_error); CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), @@ -185,7 +173,7 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), makeEdgeFeatures(), - makeTensor({1})), + makeTensor({1, 1})), std::runtime_error); } diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp index 184ee7ea..b0fccbd8 100644 --- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -1,13 +1,12 @@ +#include #include #include -#include - -#include #include #include #include #include #include +#include #include #include @@ -57,7 +56,8 @@ static std::vector contiguousStrides(const std::vector& shape) } template -static AMSTensor makeTensor(std::vector shape, const std::vector& values) +static AMSTensor makeTensor(std::vector shape, + const std::vector& values) { std::vector strides = contiguousStrides(shape); auto tensor = AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); @@ -87,8 +87,8 @@ static std::uintmax_t dtypeByteWidth(const std::string& dtype) if (dtype == "int64") { return sizeof(std::int64_t); } - CATCH_FAIL("Unsupported MGN graph diffusion tensor dtype in manifest: " - << dtype); + CATCH_FAIL( + "Unsupported MGN graph diffusion tensor dtype in manifest: " << dtype); return 0; } @@ -119,9 +119,10 @@ static TensorMetadata parseTensorMetadata(const json& tensor) return metadata; } -static void validateTensorMetadata(const TensorMetadata& metadata, - const std::string& expected_dtype, - const std::vector& expected_shape) +static void validateTensorMetadata( + const TensorMetadata& metadata, + const std::string& expected_dtype, + const std::vector& expected_shape) { // Validate the manifest before allocating AMSTensors. The binary files are // intentionally simple raw bytes, so the manifest is the only place where @@ -142,10 +143,11 @@ static void validateTensorMetadata(const TensorMetadata& metadata, } template -static std::vector readTensorBinary(const std::filesystem::path& manifest_dir, - const json& tensor, - const std::string& expected_dtype, - const std::vector& expected_shape) +static std::vector readTensorBinary( + const std::filesystem::path& manifest_dir, + const json& tensor, + const std::string& expected_dtype, + const std::vector& expected_shape) { TensorMetadata metadata = parseTensorMetadata(tensor); validateTensorMetadata(metadata, expected_dtype, expected_shape); @@ -204,11 +206,13 @@ static json loadManifest(const std::filesystem::path& fixture_dir) } static std::filesystem::path resolveManifestRelativePath( - const std::filesystem::path& manifest_dir, const json& object) + const std::filesystem::path& manifest_dir, + const json& object) { // Model and tensor paths in the manifest are relative by design. Avoiding // absolute paths makes fixtures reusable if the whole build directory moves. - const std::filesystem::path relative_path = object.at("path").get(); + const std::filesystem::path relative_path = + object.at("path").get(); CATCH_REQUIRE_FALSE(relative_path.empty()); CATCH_REQUIRE_FALSE(relative_path.is_absolute()); for (const auto& part : relative_path) { @@ -236,14 +240,22 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, CATCH_REQUIRE(global_dim == kGlobalFeatureDim); const json& tensors = graph_case.at("tensors"); - auto node_features = readTensorBinary( - manifest_dir, tensors.at("node_features"), "float32", {num_nodes, node_dim}); - auto edge_index = readTensorBinary( - manifest_dir, tensors.at("edge_index"), "int64", {2, num_edges}); - auto edge_features = readTensorBinary( - manifest_dir, tensors.at("edge_features"), "float32", {num_edges, edge_dim}); - auto global_features = readTensorBinary( - manifest_dir, tensors.at("global_features"), "float32", {1, global_dim}); + auto node_features = readTensorBinary(manifest_dir, + tensors.at("node_features"), + "float32", + {num_nodes, node_dim}); + auto edge_index = readTensorBinary(manifest_dir, + tensors.at("edge_index"), + "int64", + {2, num_edges}); + auto edge_features = readTensorBinary(manifest_dir, + tensors.at("edge_features"), + "float32", + {num_edges, edge_dim}); + auto global_features = readTensorBinary(manifest_dir, + tensors.at("global_features"), + "float32", + {1, global_dim}); return AMSHomogeneousGraph( makeTensor({toDim(num_nodes), toDim(node_dim)}, node_features), @@ -253,14 +265,16 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, } static std::vector loadReferenceDeltaU( - const std::filesystem::path& manifest_dir, const json& graph_case) + const std::filesystem::path& manifest_dir, + const json& graph_case) { const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t output_dim = graph_case.at("reference_output_dim").get(); CATCH_REQUIRE(output_dim == kReferenceOutputDim); return readTensorBinary(manifest_dir, - graph_case.at("tensors").at("reference_delta_u"), + graph_case.at("tensors").at("reference_delta_" + "u"), "float32", {num_nodes, output_dim}); } @@ -288,9 +302,9 @@ static void verifyDeltaU(const json& graph_case, const float* actual = delta_u.data(); const std::int64_t count = num_nodes * output_dim; for (std::int64_t i = 0; i < count; ++i) { - CATCH_REQUIRE(actual[i] == Catch::Approx(reference_delta_u[i]) - .epsilon(rtol) - .margin(atol)); + CATCH_REQUIRE( + actual[i] == + Catch::Approx(reference_delta_u[i]).epsilon(rtol).margin(atol)); } } @@ -329,7 +343,8 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", for (const json& graph_case : manifest.at("cases")) { CATCH_REQUIRE(graph_case.at("num_nodes").get() == expected_node_counts.at(case_index)); - CATCH_DYNAMIC_SECTION("fixture " << graph_case.at("name").get()) + CATCH_DYNAMIC_SECTION("fixture " + << graph_case.at("name").get()) { AMSHomogeneousGraph graph = makeGraph(manifest_dir, graph_case); std::vector reference_delta_u = @@ -347,10 +362,10 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", graph_case.at("num_nodes").get(); const std::int64_t output_dim = graph_case.at("reference_output_dim").get(); - outputs.node_fields.set( - "delta_u", - makeTensor({toDim(num_nodes), toDim(output_dim)}, - reference_delta_u)); + outputs.node_fields.set("delta_u", + makeTensor({toDim(num_nodes), + toDim(output_dim)}, + reference_delta_u)); }; AMSHomogeneousGraphFields outputs; From 84d7342d489ffbd1c802ceff89b57896b8937155 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 9 Jun 2026 14:16:45 -0700 Subject: [PATCH 08/10] Add extensive documentation to diffusion MGN python script. --- .../models/generate_mgn_graph_diffusion.py | 225 ++++++++++++++---- 1 file changed, 174 insertions(+), 51 deletions(-) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 5ccdcfca..7d8d14dd 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -1,32 +1,55 @@ #!/usr/bin/env python3 -"""Train/export a tiny pure-Torch MGN-like graph diffusion surrogate. +"""Train/export a tiny pure-Torch graph diffusion surrogate. -This utility is a self-contained learned-graph smoke test for the AMS graph -surrogate path. It intentionally stays pure PyTorch: no PyG/DGL/PhysicsNeMo, -no MFEM data, and no extra graph-learning runtime dependencies. +This file is meant to be readable as a small end-to-end example, even if you +have not worked with AI surrogate models before. -The workflow is split into three explicit modes: +In this example, a "surrogate" is a learned replacement for a known +calculation. We first create synthetic graph data where the correct answer is +known exactly. Then we train a small neural network to imitate that answer. +Finally, we export the trained network and check that AMS/C++ gets the same +answer as Python/TorchScript for fixed test graphs. -* feasibility: create the model, run eager inference on two graph sizes, script - it, reload it, and verify all outputs match. This gates TorchScript and - dynamic N/E support before any training cost is paid. -* train: train the tiny model on deterministic synthetic graph diffusion data - and save only a checkpoint/state dict plus training metrics. -* fixtures: load the checkpoint, export the AMS-wrapped TorchScript model, and - write runtime fixture files consumed by the C++ parity test. +The data is a graph: + +* Nodes are sample points in a 2D square. Each node has features such as its + x/y position, a scalar state value u, and a conductivity-like value kappa. +* Edges connect nearby nodes. Each edge has a source node and a destination + node, so a message can travel "from source to destination". +* Global features describe the whole graph. Here the only global feature is + dt, the time-step size. + +The synthetic target is a diffusion update: + + delta_u_i = dt * sum_{src -> i} weight * (u_src - u_i) + +Intuitively, each node is pulled toward its neighbors. If a neighboring source +node has a larger u value, it sends a positive contribution. If it has a +smaller u value, it sends a negative contribution. This is close in spirit to +heat diffusion, but cheap and fully deterministic. + +The model is "MGN-like" (MeshGraphNet-like) because it repeats the same graph +communication pattern: -The graph contract matches AMS homogeneous graph inputs: +1. build/update an edge message from source node, destination node, edge data, + and global data; +2. add incoming edge messages onto each destination node; +3. update each node from its old state and the aggregated incoming messages. - node_features [N, 4] = x, y, u, kappa - edge_index [2, E] = source row, destination row - edge_features [E, 4] = dx, dy, distance, normalized diffusion message - global_features [1, 1] = dt +The workflow has three explicit modes: -The exact target is a deterministic diffusion-like update, -delta_u_i = dt * sum_{src -> i} weight * (u_src - u_i). Training uses that -target, while the C++ fixture reference uses the exported TorchScript model -output. That distinction is deliberate: the AMS test validates deployment -parity with Python TorchScript, not physical fidelity of the learned model. +* feasibility: prove that the model can run eagerly, be scripted, be reloaded, + and still accept different graph sizes. This checks deployability before + paying any training cost. +* train: train the tiny model on deterministic synthetic graphs, then save a + checkpoint containing learned weights and metrics. +* fixtures: load the checkpoint, export the AMS-wrapped TorchScript model, and + save fixed graph inputs plus Python TorchScript reference outputs. + +The C++ parity test compares AMS output to Python TorchScript output, not to +the exact synthetic formula. Training already checks that the model learned the +synthetic target. The C++ test asks a different question: "when AMS deploys this +TorchScript graph model, does it reproduce Python inference?" """ import argparse @@ -86,8 +109,31 @@ COMPARISON_RTOL = 2.0e-5 COMPARISON_ATOL = 2.0e-5 +# Glossary for readers new to graph surrogates: +# +# N: number of nodes in one graph. +# E: number of directed edges in one graph. +# node_features: table with one row per node. Here [x, y, u, kappa]. +# edge_index: integer connectivity table with shape [2, E]. Row 0 is source +# node id, row 1 is destination node id. +# edge_features: table with one row per edge. Here [dx, dy, distance, message]. +# global_features: graph-wide values. Here just [dt]. +# delta_u: the node-wise output we want the model to predict. +# dt: time-step size in the synthetic diffusion update. +# kappa: conductivity-like node value used to define edge weights. +# checkpoint: saved training result containing model weights and metadata. +# fixture: fixed test input plus expected output used by C++ for repeatability. +# TorchScript: PyTorch's serialized model format that LibTorch/C++ can load. + class MLP(nn.Module): + """A small trainable function used as the edge/node update rule. + + An MLP, or multilayer perceptron, is just a few linear layers with a + nonlinearity. Here it learns simple formulas from examples instead of us + hand-writing the formulas in C++. + """ + def __init__(self, input_dim: int, hidden_dim: int, output_dim: int): super().__init__() self.layers = nn.Sequential( @@ -101,7 +147,12 @@ def forward(self, x: Tensor) -> Tensor: class ProcessorBlock(nn.Module): - """One MGN-like processor step implemented with plain tensor operations.""" + """One round of graph communication. + + A MeshGraphNet-style model alternates between updating edges and updating + nodes. Think of this block as one "conversation round": edges look at their + source and destination nodes, then nodes collect all incoming edge messages. + """ def __init__(self, latent_dim: int, global_dim: int): super().__init__() @@ -155,7 +206,20 @@ def forward( class TinyGraphDiffusionMGN(nn.Module): - """Small fixed-shape MeshGraphNet-like model for the AMS graph contract.""" + """Small MeshGraphNet-like model for the AMS graph contract. + + The structure is: + + * encoder: turn raw node/edge numbers into latent vectors the model can + learn with; + * processor blocks: let neighboring nodes exchange information through + edges; + * decoder: turn the final node latent vectors into node:delta_u. + + The model is intentionally tiny and fixed to two processor blocks. The goal + is to validate the learned-model loop through AMS, not to build a large or + configurable production architecture. + """ def __init__( self, @@ -180,14 +244,24 @@ def forward(self, graph: Dict[str, Tensor]) -> Dict[str, Tensor]: edge_features = graph["edge_features"] global_features = graph["global_features"] + # Raw features such as x/y/u/kappa are low-dimensional physical-looking + # values. Encoders map them into a latent space where the trainable + # processor blocks can represent richer relationships. node_latent = self.node_encoder(node_features) edge_latent = self.edge_encoder(edge_features) + + # Two graph communication rounds are enough for this toy diffusion + # problem: edges gather source/destination context, then destination + # nodes receive aggregated incoming messages. node_latent, edge_latent = self.processor1( node_latent, edge_latent, edge_index, global_features ) node_latent, edge_latent = self.processor2( node_latent, edge_latent, edge_index, global_features ) + + # Decode one scalar per node. The "node:" prefix tells AMS this belongs + # in outputs.node_fields rather than edge or global fields. delta_u = self.node_decoder(node_latent) return {"node:delta_u": delta_u} @@ -234,14 +308,23 @@ def _unique_directed_edges(knn: Tensor) -> Tensor: def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor]: - """Create one deterministic synthetic graph and its exact diffusion target.""" + """Create one deterministic synthetic graph and its exact diffusion target. + + This function is the "data generator" for the example. It replaces a real + simulation code with a cheap formula that still has the graph ingredients we + care about: node values, edge directions, edge weights, aggregation onto + destination nodes, and a node-wise output field. + """ generator = make_generator(seed) + + # 1. Sample node locations in a unit square. In a mesh-based application + # these would come from mesh vertices or degrees of freedom. positions = torch.rand((num_nodes, 2), generator=generator, dtype=torch.float32) - # Use a smooth random field rather than independent random u values. This is - # still synthetic and cheap, but it looks more like a heat-equation state and - # keeps target scales stable enough to avoid feature/target normalization. + # 2. Assign each node a smooth scalar state u. Smooth values are easier to + # relate to heat diffusion than independent random noise: neighboring points + # tend to have related values, but still differ enough to produce diffusion. phases = torch.rand((1, 4), generator=generator, dtype=torch.float32) * (2.0 * torch.pi) amps = torch.rand((1, 4), generator=generator, dtype=torch.float32) * 0.5 + 0.5 x = positions[:, 0:1] @@ -253,11 +336,17 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor + amps[:, 3:4] * torch.cos(2.0 * torch.pi * (x - y) + phases[:, 3:4]) ) u = torch.tanh(0.5 * u_raw) + + # 3. Add a conductivity-like scalar kappa and a graph-wide time step dt. + # kappa changes how strongly values diffuse across edges; dt scales the + # final update for every node in the graph. kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 dt = torch.rand((1, 1), generator=generator, dtype=torch.float32) * 0.06 + 0.02 - # Build a small directed kNN graph. E varies with N because duplicate edges - # are removed after adding the reverse direction, so N=24 and N=73 exercise + # 4. Connect nearby nodes with a k-nearest-neighbor graph. The model sees + # directed edges, so edge_index[0, e] is the source node and edge_index[1, e] + # is the destination node for edge e. E varies with N because duplicate + # edges are removed after adding the reverse direction; that helps test # dynamic node and edge counts in TorchScript and C++. distances = torch.cdist(positions, positions) masked = distances + torch.eye(num_nodes, dtype=torch.float32) * 1.0e6 @@ -271,23 +360,25 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor conductivity = 0.5 * (kappa.index_select(0, src) + kappa.index_select(0, dst)) raw_weight = conductivity / (distance + 0.05) - # Normalize incoming weights per destination node. This is edge-weight - # normalization for a bounded synthetic target, not feature or target - # normalization embedded in the model. + # 5. Turn geometric distance and kappa into a diffusion weight. The incoming + # weights for each destination node are normalized so every node receives a + # bounded weighted average from its neighbors. This is only edge-weight + # normalization in the synthetic formula, not model feature normalization. incoming_sum = torch.zeros((num_nodes, 1), dtype=torch.float32) incoming_sum = incoming_sum.index_add(0, dst, raw_weight) weight = raw_weight / incoming_sum.index_select(0, dst).clamp_min(1.0e-8) - # The fourth edge feature is the already weighted scalar message. Including - # it keeps the learned problem small while still requiring source/destination - # indexing and destination aggregation to recover node:delta_u. + # 6. Build the tensors the model will receive. The fourth edge feature is + # the weighted scalar message weight * (u_src - u_dst). Including it keeps + # the learned problem small while still requiring source/destination indexing + # and destination aggregation to recover node:delta_u. node_features = torch.cat((positions, u, kappa), dim=1).contiguous() messages = weight * (u.index_select(0, src) - u.index_select(0, dst)) edge_features = torch.cat((delta_pos, distance, messages), dim=1).contiguous() - # Exact supervised target: aggregate incoming weighted messages and scale by - # dt. Python training checks the model against this target; C++ checks AMS - # against TorchScript reference outputs generated after export. + # 7. Compute the exact supervised answer. For each edge, the source node + # pushes the destination node toward u_src. index_add sums all incoming edge + # messages for each destination node, then dt scales the update. target = torch.zeros((num_nodes, 1), dtype=torch.float32) target = target.index_add(0, dst, messages) target = dt * target @@ -319,9 +410,18 @@ def assert_close(name: str, actual: Tensor, expected: Tensor, atol: float = 1.0e def run_feasibility(out_dir: Path) -> None: - # This mode owns only TorchScript compatibility. It intentionally does not - # train or write deployment fixtures, so failures point at scripting, - # dynamic graph sizes, or AMS wrapper compatibility. + # Mode 1: feasibility. + # + # Before training anything, prove that this model is deployable in the form + # AMS needs. We run the same randomly initialized model four ways: + # + # 1. normal eager PyTorch; + # 2. torch.jit.script output; + # 3. the scripted model after saving and loading it; + # 4. the AMS graph wrapper around the model. + # + # If these disagree, training would only hide a deployment problem. This + # mode intentionally does not write fixtures or a trained checkpoint. out_dir.mkdir(parents=True, exist_ok=True) print("[info] MGN graph diffusion feasibility gate") print(f"[info] model_seed={MODEL_SEED}, fixture_sizes={FIXTURE_GRAPH_SIZES}, config={model_config()}") @@ -386,8 +486,13 @@ def evaluate(model: nn.Module, cases: Iterable[Tuple[Dict[str, Tensor], Tensor]] def run_train(out_dir: Path) -> None: - # This mode owns learned weights and metrics only. It writes a checkpoint - # that fixtures mode must consume; fixtures mode never silently retrains. + # Mode 2: train. + # + # The model sees many freshly generated synthetic graphs. For each graph we + # know the exact target_delta_u, so training is ordinary supervised learning: + # predict delta_u, measure mean-squared error, and update the network + # weights. This mode writes learned weights and metrics only. It does not + # export the AMS model or write C++ fixture tensors. out_dir.mkdir(parents=True, exist_ok=True) print("[info] MGN graph diffusion training") print(f"[info] model_seed={MODEL_SEED}, train_data_seed={TRAIN_DATA_SEED}, val_data_seed={VAL_DATA_SEED}") @@ -412,6 +517,8 @@ def run_train(out_dir: Path) -> None: best_state_dict = copy.deepcopy(model.state_dict()) for step in range(1, TRAIN_STEPS + 1): + # Generate one graph per optimization step. Avoiding batching keeps the + # example small and makes dynamic graph sizes part of normal training. num_nodes = 16 + ((step * 53) % 113) graph, target = generate_graph(num_nodes, TRAIN_DATA_SEED + step) pred = model(graph_to_model_input(graph))["node:delta_u"] @@ -450,6 +557,9 @@ def run_train(out_dir: Path) -> None: f"[info] selected best_step={best_step}, validation_mse={val_mse:.8e}, " f"zero_baseline_mse={zero_baseline_mse:.8e}, primary_threshold={primary_threshold:.8e}" ) + # The primary success check is relative to a trivial model that always + # predicts zero. That is easier to explain and more robust to target scale + # changes than an absolute MSE alone. if val_mse > SECONDARY_ABSOLUTE_MSE: print( f"[warn] validation_mse={val_mse:.8e} is above secondary diagnostic " @@ -509,9 +619,11 @@ def load_checkpoint(path: Path) -> Dict[str, object]: def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype: str) -> Dict[str, object]: """Write one fixture tensor and return the matching manifest entry. - The C++ test deliberately avoids NumPy or torch file readers. Each tensor is - raw contiguous row-major little-endian bytes, and fixtures.json records the - dtype, shape, endianness, and byte count needed to validate it before use. + The C++ test deliberately avoids NumPy, pickle, or torch file readers for + graph inputs. Those formats would make the C++ side depend on Python data + tooling. Instead each tensor is raw contiguous row-major little-endian bytes, + and fixtures.json records the dtype, shape, endianness, and byte count + needed to validate it before use. """ rel_path = Path(relative_path) @@ -542,10 +654,13 @@ def write_tensor_binary(out_dir: Path, relative_path: str, tensor: Tensor, dtype def run_fixtures(out_dir: Path) -> None: - # This mode owns deployment artifacts: the AMS-wrapped TorchScript model, - # the fixture manifest, and raw tensor binaries. Reference outputs are model - # outputs from the reloaded TorchScript artifact, because C++ parity should - # answer "does AMS reproduce Python TorchScript inference?" + # Mode 3: fixtures. + # + # A fixture is fixed test data: inputs plus the expected output. This mode + # loads the trained checkpoint, exports the AMS-wrapped TorchScript model, + # creates two fixed graph inputs, runs the exported model in Python, and + # saves those model outputs as the C++ reference. Fixtures mode never + # silently trains because parity tests should use an explicit checkpoint. out_dir.mkdir(parents=True, exist_ok=True) checkpoint_path = out_dir / CHECKPOINT_NAME if not checkpoint_path.exists(): @@ -567,6 +682,8 @@ def run_fixtures(out_dir: Path) -> None: wrapped.save(str(model_path)) print(f"[info] Saved AMS TorchScript model: {model_path}") + # Reload the saved artifact before generating references. That catches + # export-time issues and makes the Python reference match what C++ loads. reloaded = torch.jit.load(str(model_path)).eval() cases: List[Dict[str, object]] = [] with torch.no_grad(): @@ -576,6 +693,9 @@ def run_fixtures(out_dir: Path) -> None: output = reloaded(graph)["node:delta_u"].detach().cpu().contiguous() name = f"mgn_diffusion_n{num_nodes}" case_dir = name + # Each tensor file is accompanied by manifest metadata. C++ uses the + # metadata to validate shape/type/size before constructing + # AMSTensor objects. tensors = { "node_features": write_tensor_binary( out_dir, f"{case_dir}/node_features.bin", graph["node_features"], "float32" @@ -610,6 +730,9 @@ def run_fixtures(out_dir: Path) -> None: ) cases.append(case) + # The manifest is the table of contents for the fixture directory. Paths are + # relative to fixtures.json so the generated build-tree directory can be + # moved as a unit. manifest = { "format_version": FIXTURE_FORMAT_VERSION, "endianness": FIXTURE_ENDIANNESS, From dac73a0791f9bc00eda84a6a1f1c70b3d7312d16 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 9 Jun 2026 15:44:29 -0700 Subject: [PATCH 09/10] Add extensive documentation to the C++ diffusion MGN example. --- .../test_graph_mgn_surrogate.cpp | 149 ++++++++++++++++-- 1 file changed, 138 insertions(+), 11 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp index b0fccbd8..7a52bf4d 100644 --- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -14,6 +14,25 @@ #include "AMSGraph.hpp" #include "AMSTensor.hpp" +// This test is the C++ half of the MGN diffusion example. +// +// The Python generator does the learned-model work: +// +// 1. create synthetic graph diffusion data; +// 2. train a small pure-Torch MeshGraphNet-like model; +// 3. export that model as an AMS-wrapped TorchScript file; +// 4. write "fixtures": fixed graph inputs plus Python TorchScript reference +// outputs. +// +// This C++ test does not know the diffusion formula and does not try to judge +// model quality. Its job is deployment parity: load the exact TorchScript model +// and graph tensors from the generated fixture directory, run them through AMS, +// and verify that the AMS output field node:delta_u matches the Python +// TorchScript reference output. +// +// Said differently: Python proves "the model learned something"; this test +// proves "AMS runs the exported model the same way Python does." + using namespace ams; using json = nlohmann::json; @@ -35,6 +54,25 @@ constexpr std::int64_t kNodeFeatureDim = 4; constexpr std::int64_t kEdgeFeatureDim = 4; constexpr std::int64_t kGlobalFeatureDim = 1; constexpr std::int64_t kReferenceOutputDim = 1; +constexpr double kUnusedGraphSurrogateThreshold = 1.0; +constexpr int kSingleRankProcessId = 0; +constexpr int kSingleRankWorldSize = 1; +constexpr bool kStoreTrainingData = false; +constexpr const char* kDomainName = "test_mgn_graph_diffusion_surrogate"; + +// Glossary for readers new to this path: +// +// manifest: fixtures.json, the table of contents for generated fixture files. +// fixture directory: build-tree directory containing fixtures.json, the +// TorchScript model, and raw tensor binaries. +// raw tensor binary: contiguous bytes for one tensor; shape/dtype live in JSON. +// AMSHomogeneousGraph: AMS C++ container for node_features, edge_index, +// edge_features, and optional global_features. +// node:delta_u: TorchScript output key. AMS parses this into the node field +// named "delta_u". +// fallback callback: application callback AMS calls only when surrogate +// inference is unavailable or rejected. +// parity: equality of AMS/LibTorch inference and Python TorchScript inference. struct TensorMetadata { std::filesystem::path path; @@ -46,6 +84,9 @@ struct TensorMetadata { static std::vector contiguousStrides(const std::vector& shape) { + // AMSTensor stores both shape and strides. These fixtures are written as + // simple row-major contiguous arrays, so the last dimension has stride 1 and + // each earlier stride is the product of dimensions to its right. std::vector strides(shape.size(), 1); Dim stride = 1; for (std::size_t i = shape.size(); i-- > 0;) { @@ -59,6 +100,9 @@ template static AMSTensor makeTensor(std::vector shape, const std::vector& values) { + // Convert ordinary C++ vectors read from fixture binaries into owned + // AMSTensors. From this point on, the graph looks like application-provided + // AMS input rather than test-specific storage. std::vector strides = contiguousStrides(shape); auto tensor = AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); CATCH_REQUIRE(values.size() == static_cast(tensor.elements())); @@ -75,12 +119,19 @@ static Dim toDim(std::int64_t value) static bool hostIsLittleEndian() { + // Fixture binaries are written little-endian. AMS currently reads them by + // copying bytes directly into native scalar arrays, so running this test on a + // big-endian host should fail explicitly instead of silently byte-swapping + // wrong values. const std::uint16_t value = 1; return *reinterpret_cast(&value) == 1; } static std::uintmax_t dtypeByteWidth(const std::string& dtype) { + // The generator only writes float32 tensors for real-valued graph data and + // int64 tensors for edge_index. Supporting more dtypes would require extending + // both the manifest contract and the AMSTensor construction below. if (dtype == "float32") { return sizeof(float); } @@ -104,6 +155,9 @@ static std::uintmax_t shapeElementCount(const std::vector& shape) static TensorMetadata parseTensorMetadata(const json& tensor) { + // Pull the small per-tensor manifest record into a typed C++ struct. This is + // intentionally separate from validation so missing fields, unsupported + // dtypes, and wrong shapes fail at the most helpful point in the test. CATCH_REQUIRE(tensor.contains("path")); CATCH_REQUIRE(tensor.contains("dtype")); CATCH_REQUIRE(tensor.contains("shape")); @@ -124,9 +178,13 @@ static void validateTensorMetadata( const std::string& expected_dtype, const std::vector& expected_shape) { - // Validate the manifest before allocating AMSTensors. The binary files are - // intentionally simple raw bytes, so the manifest is the only place where - // dtype, shape, byte order, and size are described. + // Validate the manifest before allocating AMSTensors. + // + // The binary files deliberately contain no headers: no magic number, no shape, + // no dtype tag. That keeps them easy for C++ to read without Python or NumPy, + // but it means fixtures.json is the source of truth. If the manifest says a + // tensor is [N, 4] float32, this check makes sure the file metadata agrees + // with exactly the tensor shape this test is about to construct. CATCH_REQUIRE_FALSE(metadata.path.empty()); CATCH_REQUIRE_FALSE(metadata.path.is_absolute()); for (const auto& part : metadata.path) { @@ -134,6 +192,9 @@ static void validateTensorMetadata( } CATCH_REQUIRE(metadata.dtype == expected_dtype); CATCH_REQUIRE(metadata.shape == expected_shape); + + // The generator writes little-endian bytes. The host-endianness check makes + // the limitation visible if this test ever runs on a different architecture. CATCH_REQUIRE(metadata.endianness == "little"); CATCH_REQUIRE(hostIsLittleEndian()); @@ -153,7 +214,8 @@ static std::vector readTensorBinary( validateTensorMetadata(metadata, expected_dtype, expected_shape); // All tensor paths are relative to fixtures.json. That keeps the generated - // fixture directory relocatable inside different build trees. + // fixture directory relocatable: moving the build directory as a unit does not + // invalidate absolute paths baked into the manifest. const std::filesystem::path tensor_path = manifest_dir / metadata.path; CATCH_REQUIRE(std::filesystem::exists(tensor_path)); CATCH_REQUIRE(std::filesystem::is_regular_file(tensor_path)); @@ -171,6 +233,9 @@ static std::vector readTensorBinary( std::numeric_limits::max())); input.read(reinterpret_cast(values.data()), static_cast(metadata.byte_size)); + // gcount verifies we actually read the promised number of bytes. Combined + // with file_size above, this catches truncated or stale fixture files before + // they can become misleading numerical comparisons. CATCH_REQUIRE(static_cast(input.gcount()) == metadata.byte_size); } @@ -179,6 +244,9 @@ static std::vector readTensorBinary( static json loadManifest(const std::filesystem::path& fixture_dir) { + // fixtures.json is generated by the Python fixtures mode. If this file is + // missing, the C++ test cannot make progress because it does not know which + // model or tensors to load. const std::filesystem::path manifest_path = fixture_dir / "fixtures.json"; if (!std::filesystem::exists(manifest_path)) { CATCH_FAIL("Missing MGN graph diffusion fixtures at " @@ -190,6 +258,7 @@ static json loadManifest(const std::filesystem::path& fixture_dir) std::ifstream input(manifest_path); CATCH_REQUIRE(input); json manifest = json::parse(input); + // Keep the manifest version check close to parsing so future fixture format // changes fail loudly instead of being interpreted as the current raw-binary // contract. @@ -225,8 +294,16 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, const json& graph_case) { // Runtime fixture binaries become the same AMSTensor-backed homogeneous graph - // that an application would pass to AMS: node features, canonical edge_index, - // edge features, and required global features. + // that an application would pass to AMS: + // + // node_features [N, 4] -> x, y, u, kappa + // edge_index [2, E] -> source row, destination row + // edge_features [E, 4] -> dx, dy, distance, message + // global_features [1, 1] -> dt + // + // Keeping this conversion here, instead of compiling arrays into the test, + // makes the test closer to a real AMS deployment path: model and tensors are + // runtime files. const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t num_edges = graph_case.at("num_edges").get(); const std::int64_t node_dim = @@ -235,11 +312,18 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, graph_case.at("edge_feature_dim").get(); const std::int64_t global_dim = graph_case.at("global_feature_dim").get(); + + // These feature dimensions are part of this particular learned example. If + // the generator changes the graph contract, the C++ parity test should fail + // and be updated deliberately. CATCH_REQUIRE(node_dim == kNodeFeatureDim); CATCH_REQUIRE(edge_dim == kEdgeFeatureDim); CATCH_REQUIRE(global_dim == kGlobalFeatureDim); const json& tensors = graph_case.at("tensors"); + + // edge_index is the only integer tensor. It must remain int64 because the + // TorchScript model canonicalizes and indexes with 64-bit node ids. auto node_features = readTensorBinary(manifest_dir, tensors.at("node_features"), "float32", @@ -268,6 +352,9 @@ static std::vector loadReferenceDeltaU( const std::filesystem::path& manifest_dir, const json& graph_case) { + // Reference output is one scalar delta_u per node. It was produced in Python + // by reloading the exported TorchScript model, so it is the closest available + // representation of what LibTorch should compute in AMS. const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t output_dim = graph_case.at("reference_output_dim").get(); @@ -286,13 +373,17 @@ static void verifyDeltaU(const json& graph_case, double atol) { // The reference is the Python TorchScript output, not the exact synthetic - // diffusion target. This test is about AMS/LibTorch deployment parity. + // diffusion target. Training mode already checks whether the model learned + // the synthetic formula. This test is about AMS/LibTorch deployment parity. const std::int64_t num_nodes = graph_case.at("num_nodes").get(); const std::int64_t output_dim = graph_case.at("reference_output_dim").get(); CATCH_REQUIRE(output_dim == kReferenceOutputDim); CATCH_REQUIRE(outputs.node_fields.contains("delta_u")); + + // The Python model returns "node:delta_u". AMS strips the "node:" namespace + // and stores the tensor under outputs.node_fields.at("delta_u"). const auto& delta_u = outputs.node_fields.at("delta_u"); CATCH_REQUIRE(delta_u.shape()[0] == num_nodes); CATCH_REQUIRE(delta_u.shape()[1] == output_dim); @@ -313,11 +404,18 @@ static void verifyDeltaU(const json& graph_case, CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", "[wf][graph][surrogate][mgn]") { + // CMake sets AMS_MGN_DIFFUSION_FIXTURE_DIR to the generated build-tree + // fixture directory. The manifest inside that directory tells us where the + // model and per-case tensor files live. const std::filesystem::path fixture_dir = AMS_MGN_DIFFUSION_FIXTURE_DIR; CATCH_REQUIRE_FALSE(fixture_dir.empty()); const json manifest = loadManifest(fixture_dir); const std::filesystem::path manifest_dir = fixture_dir; + + // The model path also comes from the manifest rather than being hard-coded in + // the test. That keeps the C++ loader tied to the same artifact table of + // contents as the tensor files. const std::filesystem::path model_path = resolveManifestRelativePath(manifest_dir, manifest.at("model")); CATCH_REQUIRE(std::filesystem::exists(model_path)); @@ -327,17 +425,40 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", const double atol = manifest.at("comparison").at("atol").get(); CATCH_REQUIRE(manifest.at("cases").size() == 2); + // Register the exported TorchScript model with AMS and create an executor + // exactly as an application would before calling AMSExecute. + // + // This test runs as a single-process executor, so process_id/world_size are + // named constants rather than bare 0/1 literals. + // + // The threshold argument has no acceptance/rejection meaning for homogeneous + // graph surrogate execution today: the graph path either runs the model or + // falls back if the model cannot be used. AMSRegisterAbstractModel still + // requires a threshold because the API is shared with pointwise surrogates. + // Use 1.0 here as a readable "accept everything" value for that unused + // graph-path parameter. + // + // The store_data argument is also inherited from the shared pointwise + // workflow API. It controls whether an AMSWorkflow opens a training-data DB, + // but this graph parity test is inference-only and requires the fallback + // callback to remain unused. Current graph fallback data storage is also not + // implemented, so false keeps the test from opening unrelated DB state. AMSInit(); const std::string model_path_string = model_path.string(); - auto model = AMSRegisterAbstractModel("test_mgn_graph_diffusion_surrogate", - 0.5, + auto model = AMSRegisterAbstractModel(kDomainName, + kUnusedGraphSurrogateThreshold, model_path_string.c_str(), - false); - AMSExecutor executor = AMSCreateExecutor(model, 0, 1); + kStoreTrainingData); + AMSExecutor executor = + AMSCreateExecutor(model, kSingleRankProcessId, kSingleRankWorldSize); // The two fixed fixture sizes intentionally have different N and E values. // Running both cases catches accidental static-shape assumptions in the // exported TorchScript model or in AMS graph tensor handling. + // + // N is checked explicitly here because this test is meant to prove the model + // handles both fixture graph sizes, not just "whatever happened to be in the + // manifest". const std::vector expected_node_counts = {24, 73}; std::size_t case_index = 0; for (const json& graph_case : manifest.at("cases")) { @@ -346,6 +467,8 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", CATCH_DYNAMIC_SECTION("fixture " << graph_case.at("name").get()) { + // Load one graph case from runtime files, then run it through AMS. Each + // case has its own N, E, input tensors, and reference output. AMSHomogeneousGraph graph = makeGraph(manifest_dir, graph_case); std::vector reference_delta_u = loadReferenceDeltaU(manifest_dir, graph_case); @@ -355,6 +478,10 @@ CATCH_TEST_CASE("AMSExecute homogeneous graph MGN diffusion surrogate", // this parity test that would hide a deployment failure, so the callback // records whether it was invoked and supplies the reference only as a // valid fallback value. + // + // The assertion below requires callback_invoked == false. That means AMS + // accepted the graph surrogate output instead of falling back to this + // lambda. HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&, AMSHomogeneousGraphFields& outputs) { callback_invoked = true; From 1294d9956c0c98483cadeb4506a9756d9d4e4113 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 9 Jun 2026 16:48:25 -0700 Subject: [PATCH 10/10] Fix MGN with new global_features rank 1. --- tests/AMSlib/ams_interface/test_graph_fallback.cpp | 7 +------ .../ams_interface/test_graph_mgn_surrogate.cpp | 6 +++--- tests/AMSlib/models/generate_mgn_graph_diffusion.py | 12 +++++++++--- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_fallback.cpp b/tests/AMSlib/ams_interface/test_graph_fallback.cpp index b261e178..5b25ef8f 100644 --- a/tests/AMSlib/ams_interface/test_graph_fallback.cpp +++ b/tests/AMSlib/ams_interface/test_graph_fallback.cpp @@ -160,11 +160,6 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", makeEdgeIndex64(), makeTensor({2, 1})), std::runtime_error); - CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), - makeEdgeIndex64(), - makeEdgeFeatures(), - makeTensor({2})), - std::runtime_error); CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), makeEdgeFeatures(), @@ -173,7 +168,7 @@ CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(), makeEdgeIndex64(), makeEdgeFeatures(), - makeTensor({1, 1})), + makeTensor({1})), std::runtime_error); } diff --git a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp index 7a52bf4d..20cb7c76 100644 --- a/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp +++ b/tests/AMSlib/ams_interface/test_graph_mgn_surrogate.cpp @@ -299,7 +299,7 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, // node_features [N, 4] -> x, y, u, kappa // edge_index [2, E] -> source row, destination row // edge_features [E, 4] -> dx, dy, distance, message - // global_features [1, 1] -> dt + // global_features [1] -> dt // // Keeping this conversion here, instead of compiling arrays into the test, // makes the test closer to a real AMS deployment path: model and tensors are @@ -339,13 +339,13 @@ static AMSHomogeneousGraph makeGraph(const std::filesystem::path& manifest_dir, auto global_features = readTensorBinary(manifest_dir, tensors.at("global_features"), "float32", - {1, global_dim}); + {global_dim}); return AMSHomogeneousGraph( makeTensor({toDim(num_nodes), toDim(node_dim)}, node_features), makeTensor({2, toDim(num_edges)}, edge_index), makeTensor({toDim(num_edges), toDim(edge_dim)}, edge_features), - makeTensor({1, toDim(global_dim)}, global_features)); + makeTensor({toDim(global_dim)}, global_features)); } static std::vector loadReferenceDeltaU( diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 7d8d14dd..2c9c9952 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -166,6 +166,12 @@ def forward( edge_index: Tensor, global_features: Tensor, ) -> Tuple[Tensor, Tensor]: + # AMS homogeneous graphs provide global_features as one graph-level + # vector with shape [F_global]. For concatenation, this block needs one + # copy per edge and one copy per node, so the broadcast happens here at + # the point where graph-wide data is turned into per-entity inputs. + global_row = global_features.unsqueeze(0) + # The canonical AMS edge_index uses row 0 as source and row 1 as # destination. index_select gathers per-edge source/destination node # states without any graph-library dependency. @@ -175,7 +181,7 @@ def forward( # Edge update: each edge sees its current latent state, source node, # destination node, and the graph-level dt feature. The residual update # keeps this tiny model easy to train. - global_edges = global_features.expand(edge_latent.shape[0], global_features.shape[1]) + global_edges = global_row.expand(edge_latent.shape[0], global_row.shape[1]) edge_input = torch.cat( ( edge_latent, @@ -199,7 +205,7 @@ def forward( # Node update: each node sees its previous latent state, the aggregated # incoming message, and the global dt feature. - global_nodes = global_features.expand(node_latent.shape[0], global_features.shape[1]) + global_nodes = global_row.expand(node_latent.shape[0], global_row.shape[1]) node_input = torch.cat((node_latent, aggregated, global_nodes), dim=1) node_latent = node_latent + self.node_mlp(node_input) return node_latent, edge_latent @@ -341,7 +347,7 @@ def generate_graph(num_nodes: int, seed: int) -> Tuple[Dict[str, Tensor], Tensor # kappa changes how strongly values diffuse across edges; dt scales the # final update for every node in the graph. kappa = torch.rand((num_nodes, 1), generator=generator, dtype=torch.float32) + 0.5 - dt = torch.rand((1, 1), generator=generator, dtype=torch.float32) * 0.06 + 0.02 + dt = torch.rand((1,), generator=generator, dtype=torch.float32) * 0.06 + 0.02 # 4. Connect nearby nodes with a k-nearest-neighbor graph. The model sees # directed edges, so edge_index[0, e] is the source node and edge_index[1, e]