Add chunked GNO forward pass - #746

import gc, json, statistics, time, sys, subprocess
from pathlib import Path
ROOT = Path.cwd()
HEAD = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=ROOT, text=True).strip()
import torch
sys.path.insert(0, str(ROOT))
from neuralop.layers.gno_block import GNOBlock
torch.manual_seed(726)
block = GNOBlock(in_channels=16, out_channels=16, coord_dim=2,
                 pos_embedding_type=None, radius=10.0, reduction="mean",
                 transform_type="nonlinear", channel_mlp_layers=[64, 64],
                 use_torch_scatter_reduce=False,
                 use_open3d_neighbor_search=False).cuda()
y = torch.rand(512, 2, device="cuda")
x = torch.rand(512, 2, device="cuda")
f = torch.rand(512, 16, device="cuda")
def measure(chunk_size, training):
    times, peaks = [], []
    for repeat in range(4):
        block.zero_grad(set_to_none=True)
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        baseline = torch.cuda.memory_allocated()
        torch.cuda.reset_peak_memory_stats()
        start = time.perf_counter()
        with torch.set_grad_enabled(training):
            output = block(y, x, f_y=f) if chunk_size is None else block.chunked_forward(y, x, f_y=f, chunk_size=chunk_size)
            if training:
                output.square().mean().backward()
        torch.cuda.synchronize()
        if repeat > 0:
            times.append((time.perf_counter() - start)*1000)
            peaks.append((torch.cuda.max_memory_allocated() - baseline)/2**20)
        del output
    return {"median_ms": statistics.median(times), "peak_extra_allocated_mib": max(peaks)}
memory_report = {
    "source_commit": HEAD,
    "torch": torch.__version__, "cuda": torch.version.cuda,
    "gpu": torch.cuda.get_device_name(0),
    "shape": {"input_points":512, "query_points":512, "edges":512*512, "features":16, "hidden":[64,64]},
    "repeats":3, "warmup":1,
    "note":"Synthetic dense graph; incremental torch-allocated memory includes graph construction. Training includes parameter gradients, not input gradients. Optional Open3D/torch-scatter disabled.",
    "results": {}
}
for training in (False, True):
    for chunk_size in (None, 4096):
        key = ("training" if training else "inference") + ("_full" if chunk_size is None else "_chunk4096")
        memory_report["results"][key] = measure(chunk_size, training)
print(json.dumps(memory_report, indent=2))
Path("gno_memory_benchmark.json").write_text(json.dumps(memory_report, indent=2))