Skip to content

Better StructArray & StaticArray support#2546

Merged
avik-pal merged 39 commits intomainfrom
ptiede-structarrays
Apr 2, 2026
Merged

Better StructArray & StaticArray support#2546
avik-pal merged 39 commits intomainfrom
ptiede-structarrays

Conversation

@ptiede
Copy link
Copy Markdown
Collaborator

@ptiede ptiede commented Feb 25, 2026

No description provided.

@ptiede ptiede changed the title Draft to figure out better StructArray support Better StructArray support Feb 25, 2026
@ptiede ptiede changed the title Better StructArray support Better StructArray & StaticArray support Feb 26, 2026
@ptiede ptiede marked this pull request as ready for review February 26, 2026 05:53
ptiede
ptiede commented Feb 26, 2026
View reviewed changes
Comment thread src/TracedUtils.jl
else
typeof(res_tmp)
end
result = similar(first(flat_args), T_res, L)
Copy link
Copy Markdown
Collaborator Author

@ptiede ptiede Feb 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be reviewed closely. I had to change this line because now flat_args may have differing array types, e.g., one could be a StructArray. Before we decided the output entirely based on the first element, which could lead to errors, e.g., 2 .* sa would try to write a StructArray to a TracedRArray.

Copy link
Copy Markdown
Collaborator

@jumerckx jumerckx Feb 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh this is clever!

@ptiede
Copy link
Copy Markdown
Collaborator Author

ptiede commented Feb 26, 2026

@wsmoses @jumerckx I think this is ready for a first review. I'll add some tests tomorrow, but this seems to enable more support for StructArrays. Right now, the only thing that could cause an issue is if one of the underlying arrays isn't a ReactantPrimitive, but that should be enough for what I need.

avik-pal
avik-pal reviewed Feb 28, 2026
View reviewed changes
Comment thread ext/ReactantStaticArraysExt.jl Outdated
wsmoses
wsmoses reviewed Mar 4, 2026
View reviewed changes
Comment thread ext/ReactantStructArraysExt.jl Outdated
sbrantq added a commit that referenced this pull request Mar 12, 2026
@sbrantq
[DROPME] squashed PR #2546
b481e46
ptiede and others added 2 commits March 23, 2026 23:22
@ptiede
Add additional changes for StaticArrays
baf5177
@ptiede @github-actions
Apply suggestions from code review
358507b 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@ptiede
Copy link
Copy Markdown
Collaborator Author

ptiede commented Mar 24, 2026

@avik-pal @Pangoraw @wsmoses I think this is ready for review. Once this is merged, most of Comrade should work.

@ptiede
Copy link
Copy Markdown
Collaborator Author

ptiede commented Mar 24, 2026

Actually one question is what are the semantics of

overloaded_mul(A, B, alpha, beta)

typically for the 5-argument mul in place is

C = alpha A*B + beta * C

I'm not sure what this should be when C isn't defined.

wsmoses
wsmoses reviewed Mar 25, 2026
View reviewed changes
Comment thread ext/ReactantStructArraysExt.jl Outdated
wsmoses
wsmoses reviewed Mar 25, 2026
View reviewed changes
Comment thread src/TracedUtils.jl Outdated
wsmoses
wsmoses reviewed Mar 25, 2026
View reviewed changes
Comment thread src/TracedUtils.jl Outdated
wsmoses
wsmoses reviewed Mar 25, 2026
View reviewed changes
Comment thread src/TracedUtils.jl Outdated
wsmoses
wsmoses reviewed Mar 25, 2026
View reviewed changes
Comment thread src/TracedUtils.jl Outdated
wsmoses
wsmoses reviewed Apr 1, 2026
View reviewed changes
Comment thread ext/ReactantStructArraysExt.jl Outdated
) where {T,TI<:Integer}
valsT = maybe_convert_elt(T, vals)
foreachfield((col, val) -> (@inbounds col[I] = val), s, valsT)
foreachfield((col, val) -> (@inbounds Reactant.@allowscalar col[I] = val), s, valsT)
Copy link
Copy Markdown
Member

@wsmoses wsmoses Apr 1, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lambda functions are going to cause some compile time issues, can we get rid of them with functions defined directly (possibly with fix1/etc as relevant)?

wsmoses
wsmoses approved these changes Apr 1, 2026
View reviewed changes
Copy link
Copy Markdown
Member

@wsmoses wsmoses left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm minus lambda nits

github-actions[bot]
github-actions Bot reviewed Apr 1, 2026
View reviewed changes
Comment thread test/integration/structarrays.jl
@avik-pal avik-pal merged commit feba0ca into main Apr 2, 2026
94 of 128 checks passed
@avik-pal avik-pal deleted the ptiede-structarrays branch April 2, 2026 04:16
Pangoraw pushed a commit that referenced this pull request Apr 3, 2026
@ptiede @wsmoses
@github-actions @avik-pal @Pangoraw
feat: better StructArray & StaticArray support (#2546)
a0a839b 
* Draft to figure out better StructArray support

* Simplify and generalize structarray type conversion

* Start adding StaticArray support

* Add StaticArray support and tweak elem_apply_while_loop to select correct container type

* Revert tracing.jl

* Remove info debug

* Remove get_ith

* Add _copyto!

* format

* Fix broken test and add new tests

* format

* add StaticArrays

* Add LinearAlgebra

* Remove unused function

* Reuse the known destination for while loop if possible

* Update ext/ReactantStructArraysExt.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Proposed improved support for SArrays

* fix dumb mistake

* Add additional changes for StaticArrays

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Cleanup

* Update

* Update

* Format

* Fix for code review

* Add comments

* So dumb

* Correct comment in overloaded_mul function

Fix comment typo in overloaded_mul function.

* Update to remove anonymous functions

* Update

* Add a complex test

* Update

---------

Co-authored-by: Billy Moses <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Avik Pal <[email protected]>
Pangoraw added a commit that referenced this pull request Apr 6, 2026
@wsmoses @Pangoraw
@github-actions @claude @SouthEndMusic @enzyme-ci-bot @maximilian-gelbrecht @avik-pal @dkytezab @ptiede
Bf16cu (#2694)
c5d087b 
* force fusion

* BFloat16 cuda kernel

* add BFloat16s dep

* only run bf16 on 1.12

* kernel cast 🪄

* depend on Enzyme-jax kernel-cast pass

* remove xla options

* Update src/Compiler.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* kernelcast

* Copy preserved argument if needed (#2722)

* Make copyto! force a new buffer

* Update src/Compiler.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* setpahts!

* yo

* Update src/TracedRArray.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* add aliasing test

* pjrt

* Update test/core/aliasing.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Move :new_buffer marking to Base.materialize! instead of _copyto!

The :new_buffer path should only be set on the broadcast path (a .= b),
not on all copyto! calls that happen to receive a Broadcasted argument.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* cleanup

* Add buffer aliasing test for struct field reassignment with both new buffers

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* less eager

* Update src/Compiler.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* use count

* fmt

* good args

* fixup

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>

* Make the CUDA.jl not loaded error more specific (#2743)

* Update LIBTPU_VERSION to 0.0.39.dev20260328

* Update Project.toml

* Update WORKSPACE

* Update ENZYMEXLA_COMMIT to a new hash

* Update ENZYMEXLA_COMMIT to new hash

* Regenerate MLIR Bindings (#2750)

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

* Bump version and update Reactant_jll dependency

* Regenerate MLIR Bindings (#2756)

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

* Multifloat options (#2757)

* Multifloat options

* fix

* fix

* fix

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update Project.toml

* `div` for ConcreteRNumber (#2674)

* div for ConcreteRNumber

* Update test/core/math_ops.jl

Co-authored-by: Paul Berg <[email protected]>

* Update src/ConcreteRArray.jl

Co-authored-by: Paul Berg <[email protected]>

* Update src/ConcreteRArray.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/ConcreteRArray.jl

Co-authored-by: Paul Berg <[email protected]>

* Update src/ConcreteRArray.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/ConcreteRArray.jl

* remove methods

* Revert "remove methods"

This reverts commit ab20e23.

* meth

* fix ambiguities

---------

Co-authored-by: Paul Berg <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Paul Berg <[email protected]>

* Multifloat fix pass parsing (#2759)

* Change shardy_passes default value to post_sdy_propagation

This is currently what the docs say

* Bump version from 0.2.246 to 0.2.247

* Update ENZYMEXLA_COMMIT hash in WORKSPACE

* Update ENZYMEXLA_COMMIT to new hash (#2760)

* Regenerate MLIR Bindings (#2761)

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

* Update LIBTPU_VERSION to 0.0.39.dev20260401

* fix: qa test (#2763)

* feat: better StructArray & StaticArray support (#2546)

* Draft to figure out better StructArray support

* Simplify and generalize structarray type conversion

* Start adding StaticArray support

* Add StaticArray support and tweak elem_apply_while_loop to select correct container type

* Revert tracing.jl

* Remove info debug

* Remove get_ith

* Add _copyto!

* format

* Fix broken test and add new tests

* format

* add StaticArrays

* Add LinearAlgebra

* Remove unused function

* Reuse the known destination for while loop if possible

* Update ext/ReactantStructArraysExt.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Proposed improved support for SArrays

* fix dumb mistake

* Add additional changes for StaticArrays

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Cleanup

* Update

* Update

* Format

* Fix for code review

* Add comments

* So dumb

* Correct comment in overloaded_mul function

Fix comment typo in overloaded_mul function.

* Update to remove anonymous functions

* Update

* Add a complex test

* Update

---------

Co-authored-by: Billy Moses <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Avik Pal <[email protected]>

* Update WORKSPACE

* untested

* turn it down

* Attach data-layout to MLIR module

* claude is cooking

* dl

* Add BFloat16 extension

* compat

* make a hard dep

---------

Co-authored-by: Paul Berg <[email protected]>
Co-authored-by: Paul Berg <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
Co-authored-by: Bart de Koning <[email protected]>
Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>
Co-authored-by: Maximilian Gelbrecht <[email protected]>
Co-authored-by: Avik Pal <[email protected]>
Co-authored-by: dkytezab <[email protected]>
Co-authored-by: Paul Tiede <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@avik-pal avik-pal avik-pal left review comments

@jumerckx jumerckx jumerckx left review comments

@github-actions github-actions[bot] github-actions[bot] left review comments

@wsmoses wsmoses wsmoses approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@ptiede @wsmoses @avik-pal @jumerckx