Repository navigation
FEM SpaceRestriction and SpacePartition rebuild do not invalidate cached Arg values #1852
Description
Bug Description
Two FEM rebuild() implementations reallocate the arrays referenced by that Arg but didn't call .invalidate():
SpaceRestriction.rebuild(): It reallocates the restriction index arrays when the domain size changes, but the cached NodeArg keeps the old pointers. Issue is reproduced.
NodePartition.rebuild() / EnvironmentSpacePartition.rebuild(): also missing partition_arg_value.invalidate(). Not reproducible because in normal same-device rebuilds, _space_to_partition keeps shape 2 * topology.node_count() and receives the new mapping in place, so the cached Arg still points to current data and no failure was reproduced. It is still worth clearing the cache here too, so both rebuild handle cached Args the same way.
Source:
Missing self.node_arg_value.invalidate(self) (reproduced):
warp/warp/_src/fem/space/restriction.py
Lines 62 to 64 in 7460a3e
Missing self.partition_arg_value.invalidate(self) (same pattern, not reproduced):
warp/warp/_src/fem/space/partition.py
Lines 168 to 178 in 7460a3e
warp/warp/_src/fem/space/partition.py
Lines 485 to 490 in 7460a3e
Repro Script:
import warp as wp
import warp.fem as fem
wp.init()
device = "cuda:0"
with wp.ScopedDevice(device):
geo = fem.Grid2D(res=wp.vec2i(4, 4))
space = fem.make_polynomial_space(geo, degree=1)
mask = wp.zeros(geo.cell_count(), dtype=int)
mask[:2].fill_(1)
geo_partition = fem.ExplicitGeometryPartition(geo, mask)
space_partition = fem.make_space_partition(
space_topology=space.topology, geometry_partition=geo_partition, with_halo=False
)
restriction = fem.make_space_restriction(
space_partition=space_partition, domain=fem.Cells(geo_partition), device=device
)
cached_before = restriction.node_arg_value(device)
node_count_before = restriction.node_count_sync()
print("node count before: ", node_count_before)
print("cached mapping before:", cached_before.dof_partition_indices.numpy()[:node_count_before])
mask[:8].fill_(1)
geo_partition.rebuild(mask)
space_partition.rebuild(device=device)
restriction.rebuild(device=device)
node_count = restriction.node_count_sync()
cached_arg = restriction.node_arg_value(device)
print("node count after: ", node_count)
print("current mapping:", restriction._dof_partition_indices.numpy()[:node_count])
print("cached mapping: ", cached_arg.dof_partition_indices.numpy()[:node_count])
assert cached_arg.dof_partition_indices.ptr == restriction._dof_partition_indices.ptr, (
"stale NodeArg after rebuild(): cached node to partition mapping does not match restriction"
)
print("NodeArg is up to date")Expected output:
node count before: 6
cached mapping before: [0 1 2 3 4 5]
node count after: 15
current mapping: [0 1 2 3 4 5 6 7 8 9 10 11 12 13 14]
cached mapping: [0 1 2 3 4 5 6 7 8 9 10 11 12 13 14]
NodeArg is up to date
Actual output:
node count before: 6
cached mapping before: [0 1 2 3 4 5]
node count after: 15
current mapping: [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14]
cached mapping: [ 0 1 2 3 4 5
2147483647 0]
Traceback (most recent call last):
File "run_p16_node_mapping.py", line 37, in <module>
assert cached_arg.dof_partition_indices.ptr == restriction._dof_partition_indices.ptr, (
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: stale NodeArg after rebuild(): cached node to partition mapping does not match restriction
System Information
Warp: 1.17.0rc1
Python: 3.12
OS: L4T R38.2
Device: NVIDIA Jetson Thor
CUDA: 13.0.0
Activity
shi-eric commentedon Aug 26, 2026
This is addressed by 808ddbd.
The change invalidates cached FEM argument structs before rebuilding
SpaceRestriction,NodePartition, andEnvironmentSpacePartition. This prevents a previously cachedNodeArgorPartitionArgfrom retaining stale array references when a rebuild reallocates internal mappings.Public API behavior
Test coverage
test_restriction_rebuildinwarp/tests/fem/test_fem_integrate.py, registered on all detected devices, expands an explicit geometry partition and verifies that the restriction's cachedNodeArgpoints to the rebuilt mapping.Behavioral verification
Issue-shaped checks passed on CPU and CUDA: the original restriction scenario refreshed its mapping from 6 to 15 nodes, and
NodePartitionandEnvironmentSpacePartitionboth returned fresh cached argument structs after rebuild.Follow-up to consider
NodePartitionandEnvironmentSpacePartitioncache invalidation. The additional issue-shaped verification passed, so this does not block closure.Invalidate cached FEM arguments after rebuilds [NVIDIAGH-1852]
Invalidate cached FEM arguments after rebuilds [NVIDIAGH-1852]