Skip to content

🚀[FEA]: ISLA follow-ups from the review of #2257 #2263

Open

Description

@peterdsharpe

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Low (would be nice)

Please provide a clear description of problem you would like to solve.

These items came up in the review of #2257, which adds physicsnemo.experimental.models.isla. Each one is valid. Each is left for later because it is larger than a review fix, or because it changes a summation order. The PR leaves a note in the code at most of these places.

  • Build the boundary geometry's search structure once per case. exact_boundary_geometry rebuilds the weld, the cell tree and the winding moments for every chunk and every forward, inside its custom ops: 4 times in one forward at query_chunk_pairs=10_000 on a 1,280-triangle sphere with 2,000 prediction points; on one GPU, a 5.2M-triangle surface with 500k points took 1.4 s in one chunk and 6.5 s in ten. The fix splits the custom ops into a setup op that returns the tree records, slabs and winding tree, and a query op that reads them, with the slab choice made independent of each chunk's point count. An argument to pass that structure across steps of the same case comes after. (review comment on model.py:2093)
  • Memory controls. Derive the default of query_chunk_pairs from free device memory, and add a bound on the surface forward's memory, which today grows with the cell count. (review comment on model.py:434)
  • Report the path each region took. Add a read-only attribute that reports the resolved implementation of each region (geometry region, tree read, geometry search), and warn on CUDA when the tree read or the geometry search falls back to its reference, as the geometry region now does. (review comment on model.py:1531)
  • _min_norm_points for many tied directions. Its cost grows as K^(d+1) in the number of tied directions (K^4 in three dimensions), for example on the axis of a faceted tube or at a vertex shared by many cells. Use a batched active-set (Wolfe) solver, and keep the result under the same optimality certificate, with the enumeration as the fallback, so the result does not change. (review comment on _geometry.py:679)
  • Bound the geometry kernel's backward partials. The anchor partials have shape (B, N / BLOCK_N, S, 3), which is large at 1,024 slices. Each program could sum several tiles and write one partial, as _basis_bwd_kernel does. This changes the backward's summation order. (review comment on _geo_kernel.py:786)
  • Second derivatives through the fused kernels. torch.autograd.grad(..., create_graph=True) through a fused kernel drops its second-order term without an error; only a second backward() raises. Second-order derivatives need reference_kernels=True today. The fast paths could refuse create_graph=True outright.
  • One Triton guard. Let physicsnemo.core.version_check.OptionalImport take a minimum version, so that HealDA and ISLA share one guard, and declare Triton's floor in pyproject.toml's optional dependencies. (review comment on _geometry.py:36)

Describe any alternatives you have considered

Each item is optional for the experimental model as it stands. Its outputs are correct today, and these items improve speed, memory, error reporting and maintenance.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    ? - Needs TriageNeed team to review and classifyenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions