Add support for indexed arrays in structs and related kernels#1328
Add support for indexed arrays in structs and related kernels#1328FabienPean-Virtonomy wants to merge 1 commit intoNVIDIA:mainfrom
Conversation
📝 WalkthroughWalkthroughAdds support for using Changes
Sequence DiagramsequenceDiagram
participant User
participant StructDef as Struct Definition
participant Setter as Field Setter / Constructor
participant Marshalling as Marshalling Layer
participant Device as Device / NumPy Backend
participant Native as Native Kernel Layer
User->>StructDef: define `@wp.struct` with `indexedarray` field
StructDef->>Setter: call _make_struct_field_constructor / _make_struct_field_setter
Setter->>Setter: recognize `indexedarray`, create constructor/setter stubs
User->>Marshalling: instance.to(device)
Marshalling->>Device: call member.value.to(device)
Device-->>Marshalling: return device-bound array/array-like
Marshalling->>Marshalling: wrap/ensure `indexedarray(...)` if needed
User->>Marshalling: instance.numpy_value() / Struct.numpy_dtype()
Marshalling->>Device: call indexedarray.numpy_value() / indexedarray_t.numpy_dtype()
Device-->>Marshalling: return NumPy-compatible structured value/dtype
Marshalling-->>User: structured numpy representation
User->>Native: launch kernel accessing indexedarray in struct
Native->>Native: atomic_add overload resolves for indexedarray
Native-->>User: kernel executes and updates indexedarray
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
FabienPean-Virtonomy
force-pushed
the
FabienPean/indexedarray_in_struct
branch
from
a9690f0 to
06e268d
Compare
Greptile SummaryThis PR extends Warp's Key observations:
Confidence Score: 5/5Safe to merge; all remaining findings are P2 style/documentation suggestions with no blocking correctness issues. The implementation correctly handles all struct code paths for
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["StructInstance.to(device)"] --> B{"field type?"}
B -- "array" --> C["value.to(device) → warp.array\nsetattr(dst, name, result)"]
B -- "indexedarray" --> D["cloned = value.to(device)"]
D --> E{"isinstance cloned indexedarray?"}
E -- "True: same device, .to returned self" --> F["setattr dst name cloned\noriginal indexedarray preserved"]
E -- "False: cross-device, clone returned warp.array" --> G["indexedarray wrap cloned\nindices become None"]
G --> H["setattr dst name wrapped"]
B -- "Struct" --> I["Recurse into nested struct fields"]
B -- "other" --> J["setattr dst name value"]
subgraph "set_indexedarray_value setter"
K{"value is None?"}
K -- Yes --> L["null indexedarray_t via var_type.__ctype__()"]
K -- No --> M["validate dtype then set __ctype__()"]
end
Reviews (4): Last reviewed commit: "Add support for indexed arrays in struct..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
warp/tests/test_indexedarray.py (1)
92-158: Add a regression test for nullable struct fields +.to(device)The linked objective includes assigning
Noneto indexedarray struct fields and moving structs across devices. Adding one focused test fors.iarr = Nonefollowed bys.to(device)would lock this behavior and prevent regressions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@warp/tests/test_indexedarray.py` around lines 92 - 158, Add a regression test that assigns None to an IndexedArrayStruct field and then calls .to(device) to ensure nullable struct fields survive device moves: create an instance of IndexedArrayStruct, set s.iarr = None, call s.to(device) (or struct_arr.to(device) if testing arrays), and assert the field remains None (or numpy_dtype/ numpy_value reflect null) and no exceptions are raised; reference IndexedArrayStruct, s.iarr, and the .to(device) call in the new test (e.g., test_indexedarray_nullable_field_to_device).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@CHANGELOG.md`:
- Line 80: Update the changelog entry for the fix "Fix `@wp.struct` packing to
support `wp.indexedarray` fields (assignment, device transfer, and NumPy
structured values)" to include the issue reference that this PR closes (GH-1327)
by appending the issue link reference format (e.g.
"([GH-1327](https://github.com/<org>/<repo>/issues/1327))") to that line so it
complies with the changelog guideline; ensure the exact entry text remains
unchanged except for adding the issue reference.
In `@warp/_src/codegen.py`:
- Around line 333-335: StructInstance.to() currently calls value.to(device) for
fields detected by matches_array_class(var.type, indexedarray) without checking
for None; change that branch to guard nullable indexedarray fields by checking
if value is not None before calling value.to(device) and otherwise set the
destination attribute to None (i.e., replace the unconditional setattr(dst,
name, value.to(device)) with a conditional that calls value.to(device) only when
value is not None and sets dst.name to None when it is).
In `@warp/_src/types.py`:
- Around line 2207-2218: The numpy_value method currently forwards
self.data.numpy_value(), but when self.data is None the fallback ends up
creating a 1-D array which loses the original ndim; update numpy_value so that
if self.data is None you construct a placeholder numpy/ctype array with ndim
equal to len(self.shape) (or the stored ndim if available) and the same shape
(use self.shape to derive dimensions) before calling its numpy_value(), so
serialized tuples preserve the original ndim for wp.indexedarray fields; keep
the existing indices handling and return signature (self.data.numpy_value(),
indices, list(self.shape)).
In `@warp/tests/test_indexedarray.py`:
- Line 104: After each wp.launch call in the test functions (e.g., the
wp.launch(kernel_indexedarray_in_struct, dim=iarr.size, inputs=[s,
expected_arr], device=device) invocation), add an explicit
wp.synchronize_device(device) immediately after the launch so kernel-side
assertions (wp.expect_eq) are flushed and failures surface deterministically;
apply the same change for the other wp.launch calls referenced (the launches at
the other test locations around lines 122 and 138) to ensure all launches are
followed by wp.synchronize_device(device).
---
Nitpick comments:
In `@warp/tests/test_indexedarray.py`:
- Around line 92-158: Add a regression test that assigns None to an
IndexedArrayStruct field and then calls .to(device) to ensure nullable struct
fields survive device moves: create an instance of IndexedArrayStruct, set
s.iarr = None, call s.to(device) (or struct_arr.to(device) if testing arrays),
and assert the field remains None (or numpy_dtype/ numpy_value reflect null) and
no exceptions are raised; reference IndexedArrayStruct, s.iarr, and the
.to(device) call in the new test (e.g.,
test_indexedarray_nullable_field_to_device).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: b9dd14c8-7838-40c4-8118-767015ef7d55
📒 Files selected for processing (5)
CHANGELOG.mdwarp/_src/codegen.pywarp/_src/types.pywarp/native/array.hwarp/tests/test_indexedarray.py
warp/_src/codegen.py
Outdated
| elif matches_array_class(var.type, indexedarray): | ||
| # indexedarray_t | ||
| setattr(dst, name, value.to(device)) |
There was a problem hiding this comment.
Guard nullable indexedarray fields before device transfer
StructInstance.to() currently calls value.to(device) unconditionally for indexedarray fields. If the field is set to None (which your setter supports), this raises at runtime.
💡 Proposed fix
- if matches_array_class(var.type, array):
- # array_t
- setattr(dst, name, value.to(device))
- elif matches_array_class(var.type, indexedarray):
- # indexedarray_t
- setattr(dst, name, value.to(device))
+ if matches_array_class(var.type, array) or matches_array_class(var.type, indexedarray):
+ # array_t / indexedarray_t
+ setattr(dst, name, None if value is None else value.to(device))🧰 Tools
🪛 Ruff (0.15.7)
[error] 333-333: matches_array_class may be undefined, or defined from star imports
(F405)
[error] 333-333: indexedarray may be undefined, or defined from star imports
(F405)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@warp/_src/codegen.py` around lines 333 - 335, StructInstance.to() currently
calls value.to(device) for fields detected by matches_array_class(var.type,
indexedarray) without checking for None; change that branch to guard nullable
indexedarray fields by checking if value is not None before calling
value.to(device) and otherwise set the destination attribute to None (i.e.,
replace the unconditional setattr(dst, name, value.to(device)) with a
conditional that calls value.to(device) only when value is not None and sets
dst.name to None when it is).
| def numpy_value(self): | ||
| # pointers are represented as unsigned 64-bit integers | ||
| indices = [] | ||
| for i in range(ARRAY_MAX_DIMS): | ||
| v = self.indices[i] | ||
| # v may be a ctypes.c_void_p instance | ||
| if isinstance(v, ctypes.c_void_p): | ||
| indices.append(0 if v.value is None else int(v.value)) | ||
| else: | ||
| indices.append(0 if v is None else int(v)) | ||
|
|
||
| return (self.data.numpy_value(), indices, list(self.shape)) |
There was a problem hiding this comment.
Preserve ndim when serializing None indexedarrays.
numpy_value() forwards self.data.numpy_value(), but the data is None branch still seeds self.data via array().__ctype__(), so every null indexedarray is serialized with data.ndim == 1. That breaks wp.indexedarray(..., ndim>1) struct fields when they are set to None and marshalled through the new structured NumPy path.
💡 Proposed fix
class indexedarray_t(ctypes.Structure):
def __init__(self, data, indices, shape):
if data is None:
- self.data = array().__ctype__()
+ ndim = len(shape)
+ self.data = array_t(data=0, grad=0, ndim=ndim, shape=(0,) * ndim, strides=(0,) * ndim)
for i in range(ARRAY_MAX_DIMS):
self.indices[i] = ctypes.c_void_p(None)
self.shape[i] = 0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@warp/_src/types.py` around lines 2207 - 2218, The numpy_value method
currently forwards self.data.numpy_value(), but when self.data is None the
fallback ends up creating a 1-D array which loses the original ndim; update
numpy_value so that if self.data is None you construct a placeholder numpy/ctype
array with ndim equal to len(self.shape) (or the stored ndim if available) and
the same shape (use self.shape to derive dimensions) before calling its
numpy_value(), so serialized tuples preserve the original ndim for
wp.indexedarray fields; keep the existing indices handling and return signature
(self.data.numpy_value(), indices, list(self.shape)).
nvlukasz
left a comment
nvlukasz
left a comment
There was a problem hiding this comment.
Thanks, this looks pretty good. Just a few minor comments worth addressing.
warp/_src/codegen.py
Outdated
| setattr(dst, name, value.to(device)) | ||
| elif matches_array_class(var.type, indexedarray): | ||
| # indexedarray_t | ||
| setattr(dst, name, value.to(device)) |
There was a problem hiding this comment.
Note that indexedarray.to() will create a regular array if it needs to be copied to a different device.
We could consider doing a full clone of data + indices to create an indexedarray, but I think that's beyond the scope of this PR. So in the meantime, we could raise an exception if the device does not match.
There was a problem hiding this comment.
Ah indeed, and the test only test same device copy. I find the shrinking to an array a good approach in general, I would prefer to leave it up to the user to consider where they store the backing array if they care about reuse or optimization.
An easy way around would be to use the passthrough property of indexedarray_t when indices=None as it would behave like an array_t if I read right.
So it means modifying either at codegen point above, or adding method to on python side indexedarray, or modifying noncontiguous base to method. Would any of the three be ok with you, are there issues I did not foresee?
There was a problem hiding this comment.
The different possibilities:
Option 1: at codegen, coerce to indexedarray only in struct packing
setattr(dst, name, value.to(device) if value.device is device else wp.indexedarray(value.to(device))Option 2: in indexedarray, always return an indexedarray and discards return array path, which means a pointer check for index access compared to array path
def to(self, device):
device = warp.get_device(device)
if self.device == device:
return self
dense = warp.clone(self, device=device)
return indexedarray(data=dense, indices=None)Option 3: in noncontiguousbase, not applicable due to other child class indexedfabricarray
# types.py L4437
return type(self)(warp.clone(self, device=device))
There was a problem hiding this comment.
♻️ Duplicate comments (1)
warp/_src/codegen.py (1)
333-344:⚠️ Potential issue | 🔴 CriticalGuard nullable
indexedarrayfields before dereferencingvalue.device.Line 338 assumes every field holds a concrete object, but Lines 442-453 explicitly allow
None.s.iarr = None; s.to(device)will currently fail withAttributeErrorinstead of preserving the nullable field.💡 Suggested fix
elif matches_array_class(var.type, indexedarray): # indexedarray_t + if value is None: + setattr(dst, name, None) + continue # only support same-device no-op transfers for now # cross-device transfers `iarr.to` returns an array_t target_device = warp.get_device(device) if value.device != target_device: raise RuntimeError(Also applies to: 442-453
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@warp/_src/codegen.py` around lines 333 - 344, The nullable indexedarray field case must be guarded before dereferencing value.device: in the block that handles matches_array_class(var.type, indexedarray) check if value is None (or uses a safe attribute access) and if so set dst's field to None (preserving the nullable) instead of accessing value.device; otherwise proceed to get target_device, compare devices and call value.to(target_device) and setattr(dst, name, ...) as before. Ensure the same null-check pattern is applied to the duplicate block around lines 442-453 that also handles indexedarray fields.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@warp/_src/codegen.py`:
- Around line 333-344: The nullable indexedarray field case must be guarded
before dereferencing value.device: in the block that handles
matches_array_class(var.type, indexedarray) check if value is None (or uses a
safe attribute access) and if so set dst's field to None (preserving the
nullable) instead of accessing value.device; otherwise proceed to get
target_device, compare devices and call value.to(target_device) and setattr(dst,
name, ...) as before. Ensure the same null-check pattern is applied to the
duplicate block around lines 442-453 that also handles indexedarray fields.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 57fc8d8b-a7b8-4d73-97cb-38576bf9bf32
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/codegen.pywarp/tests/test_indexedarray.py
✅ Files skipped from review due to trivial changes (1)
- CHANGELOG.md
FabienPean-Virtonomy
force-pushed
the
FabienPean/indexedarray_in_struct
branch
from
f26c05b to
5cf0b56
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
warp/_src/types.py (1)
2206-2218:⚠️ Potential issue | 🟠 MajorPreserve
ndimwhen serializing null indexed arrays.
numpy_value()still forwardsself.data.numpy_value(), but the null-data path inindexedarray_t.__init__()seedsself.datawitharray().__ctype__(), which hard-codesndim == 1. Awp.indexedarray(..., ndim>1)struct field set toNonewill therefore round-trip through the new structuredNumPypath with the wrong dimensionality.💡 Proposed fix
class indexedarray_t(ctypes.Structure): def __init__(self, data, indices, shape): if data is None: - self.data = array().__ctype__() + ndim = len(shape) + self.data = array_t(data=0, grad=0, ndim=ndim, shape=(0,) * ndim, strides=(0,) * ndim) for i in range(ARRAY_MAX_DIMS): self.indices[i] = ctypes.c_void_p(None) self.shape[i] = 0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@warp/_src/types.py` around lines 2206 - 2218, The numpy_value() serializer for indexedarray_t should preserve the original ndim when self.data was seeded with a default placeholder (array().__ctype__()) in indexedarray_t.__init__; currently it forwards self.data.numpy_value() which hard-codes ndim==1 for null-data cases. Update numpy_value() to detect the placeholder/null-data case (the default array ctype created in indexedarray_t.__init__) and synthesize a numpy_value for an empty array that uses self.ndim (or otherwise set the correct ndim) instead of blindly using self.data.numpy_value(); keep the rest of the returned tuple (indices, shape) intact so a None-valued indexedarray(..., ndim>1) round-trips with the correct dimensionality.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@warp/_src/types.py`:
- Around line 2206-2218: The numpy_value() serializer for indexedarray_t should
preserve the original ndim when self.data was seeded with a default placeholder
(array().__ctype__()) in indexedarray_t.__init__; currently it forwards
self.data.numpy_value() which hard-codes ndim==1 for null-data cases. Update
numpy_value() to detect the placeholder/null-data case (the default array ctype
created in indexedarray_t.__init__) and synthesize a numpy_value for an empty
array that uses self.ndim (or otherwise set the correct ndim) instead of blindly
using self.data.numpy_value(); keep the rest of the returned tuple (indices,
shape) intact so a None-valued indexedarray(..., ndim>1) round-trips with the
correct dimensionality.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 227bd925-f201-4056-943e-f9a3667833a1
📒 Files selected for processing (5)
CHANGELOG.mdwarp/_src/codegen.pywarp/_src/types.pywarp/native/array.hwarp/tests/test_indexedarray.py
✅ Files skipped from review due to trivial changes (1)
- CHANGELOG.md
🚧 Files skipped from review as they are similar to previous changes (1)
- warp/native/array.h
👮 Files not reviewed due to content moderation or server errors (2)
- warp/_src/codegen.py
- warp/tests/test_indexedarray.py
Signed-off-by: Fabien Péan <pean@virtonomy.io>
FabienPean-Virtonomy
force-pushed
the
FabienPean/indexedarray_in_struct
branch
from
5cf0b56 to
02d0e90
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
warp/_src/codegen.py (1)
320-324:⚠️ Potential issue | 🔴 CriticalHandle
Nonebefore moving indexedarray fields.Line 323 still calls
value.to(device)unconditionally. The new constructor/setter path allows indexedarray fields to beNone, soStructInstance.to()now raises for a valid struct state.🐛 Proposed fix
elif matches_array_class(var.type, indexedarray): # indexedarray_t - # `.to` returns an array if on different device, force to identity indexedarray - cloned = value.to(device) - setattr(dst, name, cloned if isinstance(cloned, indexedarray) else indexedarray(cloned)) + if value is None: + setattr(dst, name, None) + else: + # `.to()` returns an array if on different device, force to identity indexedarray + cloned = value.to(device) + setattr(dst, name, cloned if isinstance(cloned, indexedarray) else indexedarray(cloned))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@warp/_src/codegen.py` around lines 320 - 324, The code calls value.to(device) unconditionally for fields matched by matches_array_class(var.type, indexedarray), but fields may be None; update the block in codegen.py (the branch handling indexedarray) to first check if value is None and if so set setattr(dst, name, None) and skip the .to() call, otherwise perform cloned = value.to(device) and set setattr(dst, name, cloned if isinstance(cloned, indexedarray) else indexedarray(cloned)); this prevents StructInstance.to() from raising when indexedarray fields are None.
🧹 Nitpick comments (1)
warp/tests/test_indexedarray.py (1)
154-160: Exercise the structured NumPy round-trip, not just the wrapper shape.This still passes if
numpy_dtype()returns incompatibleformats/offsetsor ifnumpy_value()serializes the wrong payload. Materializing a structured NumPy array from both values will catch the interop breakage this test is meant to cover.🧪 Proposed enhancement
# Just ensure these are functional for structs embedding indexedarray_t dtype = IndexedArrayStruct.numpy_dtype() value = s.numpy_value() + structured = np.array([value], dtype=np.dtype(dtype)) test.assertIsInstance(dtype, dict) test.assertEqual(dtype["names"], ["iarr"]) test.assertEqual(len(value), 1) + test.assertEqual(structured.dtype.names, ("iarr",))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@warp/tests/test_indexedarray.py` around lines 154 - 160, The test currently only checks dtype names and wrapper shape; instead materialize a structured NumPy array using the dtype returned by IndexedArrayStruct.numpy_dtype() and the value returned by s.numpy_value() to verify they interoperate correctly: call numpy.array(value, dtype=dtype) (or create a zero-length structured array and set the field) and then assert the resulting array has the expected shape and contents (e.g., len == 1 and the field "iarr" yields the same data as s.numpy_value()); update assertions to validate that the constructed structured array round-trips rather than only checking dtype["names"] and len(value).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@warp/_src/codegen.py`:
- Around line 422-433: set_indexedarray_value is missing the grad-buffer
keepalive that set_array_value uses; mirror the same logic from set_array_value:
after writing the ctype payload (the setattr(inst._ctype, field, ...)) store or
clear the strong reference to the source grad buffer in the instance's
grad-keepalive storage (the same attribute set_array_value uses, e.g.
inst._grad_keepalive[field] = value._grad_keepalive or similar) when value is
not None, and remove/clear that entry when value is None, then leave the
cls.__setattr__(inst, field, value) call as-is.
---
Duplicate comments:
In `@warp/_src/codegen.py`:
- Around line 320-324: The code calls value.to(device) unconditionally for
fields matched by matches_array_class(var.type, indexedarray), but fields may be
None; update the block in codegen.py (the branch handling indexedarray) to first
check if value is None and if so set setattr(dst, name, None) and skip the .to()
call, otherwise perform cloned = value.to(device) and set setattr(dst, name,
cloned if isinstance(cloned, indexedarray) else indexedarray(cloned)); this
prevents StructInstance.to() from raising when indexedarray fields are None.
---
Nitpick comments:
In `@warp/tests/test_indexedarray.py`:
- Around line 154-160: The test currently only checks dtype names and wrapper
shape; instead materialize a structured NumPy array using the dtype returned by
IndexedArrayStruct.numpy_dtype() and the value returned by s.numpy_value() to
verify they interoperate correctly: call numpy.array(value, dtype=dtype) (or
create a zero-length structured array and set the field) and then assert the
resulting array has the expected shape and contents (e.g., len == 1 and the
field "iarr" yields the same data as s.numpy_value()); update assertions to
validate that the constructed structured array round-trips rather than only
checking dtype["names"] and len(value).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 2cbe8274-d019-4e67-9fe0-e61a048ab4d4
📒 Files selected for processing (5)
CHANGELOG.mdwarp/_src/codegen.pywarp/_src/types.pywarp/native/array.hwarp/tests/test_indexedarray.py
✅ Files skipped from review due to trivial changes (1)
- CHANGELOG.md
🚧 Files skipped from review as they are similar to previous changes (2)
- warp/native/array.h
- warp/_src/types.py
| def set_indexedarray_value(inst, value): | ||
| if value is None: | ||
| # create indexedarray with null pointers | ||
| setattr(inst._ctype, field, var_type.__ctype__()) | ||
| else: | ||
| assert isinstance(value, indexedarray) | ||
| assert types_equal(value.dtype, var_type.dtype), ( | ||
| f"assign to struct member variable {field} failed, expected type {type_repr(var_type.dtype)}, got type {type_repr(value.dtype)}" | ||
| ) | ||
| setattr(inst._ctype, field, value.__ctype__()) | ||
|
|
||
| cls.__setattr__(inst, field, value) |
There was a problem hiding this comment.
Mirror the array-field grad keepalive here too.
set_array_value() already keeps a strong reference to the grad buffer after writing the ctype payload, but set_indexedarray_value() does not. That means an indexedarray struct field can lose its grad-buffer keepalive even though the stored indexedarray_t still carries the pointer.
🧰 Tools
🪛 Ruff (0.15.7)
[error] 427-427: indexedarray may be undefined, or defined from star imports
(F405)
[error] 428-428: types_equal may be undefined, or defined from star imports
(F405)
[error] 429-429: type_repr may be undefined, or defined from star imports
(F405)
[error] 429-429: type_repr may be undefined, or defined from star imports
(F405)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@warp/_src/codegen.py` around lines 422 - 433, set_indexedarray_value is
missing the grad-buffer keepalive that set_array_value uses; mirror the same
logic from set_array_value: after writing the ctype payload (the
setattr(inst._ctype, field, ...)) store or clear the strong reference to the
source grad buffer in the instance's grad-keepalive storage (the same attribute
set_array_value uses, e.g. inst._grad_keepalive[field] = value._grad_keepalive
or similar) when value is not None, and remove/clear that entry when value is
None, then leave the cls.__setattr__(inst, field, value) call as-is.
3 participants
Description
This PR adds support for
indexedarraywithinstruct.Closes #1327
Checklist
Test plan
Tests were added in
warp/tests/test_indexedarray.pyNew feature / enhancement
This now works
Summary by CodeRabbit
Bug Fixes
Tests
Documentation