Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: overload mul! #230

Merged
merged 5 commits into from
Nov 6, 2024
Merged

feat: overload mul! #230

merged 5 commits into from
Nov 6, 2024

Conversation

avik-pal
Copy link
Collaborator

@avik-pal avik-pal commented Nov 6, 2024

partially addresses LuxDL/Lux.jl#1025

@nospecialize(C::TracedRArray{T1,2}),
@nospecialize(A::TracedRArray{T2,2}),
@nospecialize(B::TracedRArray{T3,2}),
) where {T1,T2,T3}
Copy link
Member

Choose a reason for hiding this comment

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

Is it worth overriding the version with alpha and beta?

Copy link
Member

Choose a reason for hiding this comment

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

this is ine too tho, obviously

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

yeah we need to add that, else it does the for loop

src/TracedRArray.jl Outdated Show resolved Hide resolved
@avik-pal avik-pal requested a review from wsmoses November 6, 2024 02:58
src/TracedRArray.jl Outdated Show resolved Hide resolved
if isone(α)
C.mlir_data = res
else
C.mlir_data = (TracedRArray{T1,2}((), res, size(C)) .* α).mlir_data
Copy link
Member

Choose a reason for hiding this comment

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

doing a broadcast here probably won't optimize well for now (or rather we need to write optimizations on the batch op. Being able to driectly emit the elementwise mul of a broadcast will likely be a nontrivial perf win as a result

Copy link
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, but probably worth not using broadcast internally for perf reasons below

@wsmoses wsmoses merged commit 7492957 into main Nov 6, 2024
14 of 30 checks passed
@wsmoses wsmoses deleted the ap/inplace_mul! branch November 6, 2024 05:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants