fix(tile): reassign shared tiles to register [GH-1440]

* fix(tile): reassign shared tiles to register

Generated CUDA tile kernels can emit wp::assign() when an existing
register-backed tile variable is reassigned from a same-shape shared-backed
tile. Add the missing assignment and adjoint paths so the forward copy
succeeds, source gradients accumulate into the shared tile, and overwritten
register destination gradients are cleared.

(GH-1440)

Signed-off-by: Shivanjan Chakravorty <shivanjanc@nvidia.com>

Approved-by: Zach Corse <zcorse@nvidia.com>

See merge request omniverse/warp!2367

Author: Glitchfix

CHANGELOG.md

@@ -125,6 +125,8 @@
   now well-defined (and gradient-stable) for near-parallel edges, where it
   could previously return geometrically valid but unstable barycentric
   weights ([GH-1437](https://github.com/NVIDIA/warp/issues/1437)).
+- Fix tile reassignment from shared storage to register storage
+  ([GH-1440](https://github.com/NVIDIA/warp/issues/1440)).
 
 ### Documentation
 

warp/native/tile.h

@@ -5602,6 +5602,34 @@ inline CUDA_CALLABLE void adj_assign(
     WP_TILE_SYNC();
 }
 
+template <typename T, typename SharedLayout, bool Owner>
+inline CUDA_CALLABLE void assign(
+    tile_register_t<T, tile_layout_register_t<typename SharedLayout::Shape>>& dest,
+    const tile_shared_t<T, SharedLayout, Owner>& src
+)
+{
+    dest.assign(src.copy_to_register());
+}
+
+template <typename T, typename SharedLayout, bool Owner>
+inline CUDA_CALLABLE void adj_assign(
+    tile_register_t<T, tile_layout_register_t<typename SharedLayout::Shape>>& dest,
+    const tile_shared_t<T, SharedLayout, Owner>& src,
+    tile_register_t<T, tile_layout_register_t<typename SharedLayout::Shape>>& adj_dest,
+    tile_shared_t<T, SharedLayout, Owner>& adj_src
+)
+{
+    (void)dest;
+    (void)src;
+
+    if (adj_src.grad.ptr != nullptr) {
+        adj_src.grad_add(adj_dest);
+    }
+
+    // Overwritten destinations do not contribute to the pre-assignment dest value.
+    adj_dest.zero();
+}
+
 
 template <typename TileA, typename Scalar> inline CUDA_CALLABLE void assign(TileA& dest, int i, const Scalar& src)
 {

warp/tests/tile/test_tile_shared_memory.py

@@ -600,6 +600,58 @@ def compute(
     np.testing.assert_allclose(out.numpy(), inp_np)
 
 
+def test_tile_register_from_shared_reassign(test, device):
+    TILE_SIZE = 8
+    BLOCK_DIM = 64
+
+    @wp.kernel(module="unique")
+    def compute(
+        src: wp.array[float],
+        overwritten: wp.array[float],
+        reassigned: wp.array[float],
+        direct: wp.array[float],
+        iters: int,
+    ):
+        t = wp.tile_load(overwritten, shape=TILE_SIZE, offset=0, storage="register")
+        s = wp.tile_load(src, shape=TILE_SIZE, offset=0, storage="shared")
+
+        for _ in range(iters):
+            t = s
+
+        wp.tile_store(reassigned, t, offset=0)
+        wp.tile_store(direct, s, offset=0)
+
+    src_np = np.arange(TILE_SIZE, dtype=np.float32) + 1.0
+    overwritten_np = np.arange(TILE_SIZE, dtype=np.float32) + 101.0
+
+    src = wp.array(src_np, requires_grad=True, device=device)
+    overwritten = wp.array(overwritten_np, requires_grad=True, device=device)
+    reassigned = wp.zeros(TILE_SIZE, dtype=float, requires_grad=True, device=device)
+    direct = wp.zeros(TILE_SIZE, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            compute,
+            dim=[1],
+            inputs=[src, overwritten, reassigned, direct, 2],
+            block_dim=BLOCK_DIM,
+            device=device,
+        )
+
+    np.testing.assert_allclose(reassigned.numpy(), src_np)
+    np.testing.assert_allclose(direct.numpy(), src_np)
+
+    tape.backward(
+        grads={
+            reassigned: wp.ones_like(reassigned, device=device),
+            direct: wp.ones_like(direct, device=device),
+        }
+    )
+
+    np.testing.assert_allclose(src.grad.numpy(), np.full(TILE_SIZE, 2.0, dtype=np.float32))
+    np.testing.assert_allclose(overwritten.grad.numpy(), np.zeros(TILE_SIZE, dtype=np.float32))
+
+
 def test_tile_scatter_masked_basic(test, device):
     """Each thread writes its index; verify all values are visible after the call."""
     TILE_SIZE = 64
@@ -889,6 +941,12 @@ class TestTileSharedMemory(unittest.TestCase):
     test_tile_shared_coalesced_mat44,
     devices=devices,
 )
+add_function_test(
+    TestTileSharedMemory,
+    "test_tile_register_from_shared_reassign",
+    test_tile_register_from_shared_reassign,
+    devices=devices,
+)
 add_function_test(
     TestTileSharedMemory, "test_tile_scatter_masked_basic", test_tile_scatter_masked_basic, devices=devices
 )