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

Generalize broadcast #35

Merged
merged 9 commits into from
Jul 14, 2024
Merged

Generalize broadcast #35

merged 9 commits into from
Jul 14, 2024

Conversation

wsmoses
Copy link
Member

@wsmoses wsmoses commented Jun 29, 2024

No description provided.

@wsmoses wsmoses requested review from mofeing and avik-pal June 30, 2024 14:20
@wsmoses
Copy link
Member Author

wsmoses commented Jun 30, 2024

Companion PR to EnzymeAD/Enzyme#1952

Untested, but should enable arbitrary broadcast support (if all scalar functions are overloaded)

src/Reactant.jl Outdated
@@ -429,7 +429,7 @@ function append_path(path, i)
return (path..., i)
end

@inline function make_tracer(seen::IdDict, prev::RT, path, mode, data) where {RT}
@inline function make_tracer(seen::IdDict, prev::RT, path, mode, data; toscalar=false, tobatch=nothing) where {RT}
Copy link
Collaborator

Choose a reason for hiding this comment

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

What do these kwargs do?

Copy link
Member Author

Choose a reason for hiding this comment

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

whether the tracer we create should automatically change all arrays to scalars, or scalars to arrays

Copy link
Collaborator

@mofeing mofeing left a comment

Choose a reason for hiding this comment

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

Just some fix requests and a couple of aesthetic suggestions.

Also, we should add a test for this.

src/overloads.jl Outdated
@@ -457,6 +457,95 @@ for (jlop, hloop) in (
end
end


function elem_apply(f, args::VarArgs{Nargs})
Copy link
Collaborator

Choose a reason for hiding this comment

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

Need where here:

Suggested change
function elem_apply(f, args::VarArgs{Nargs})
function elem_apply(f, args::VarArgs{Nargs}) where {Nargs}

src/utils.jl Outdated
@@ -13,7 +13,7 @@ function apply(f, args...; kwargs...)
return f(args...; kwargs...)
end

function make_mlir_fn(mod, f, args, kwargs, name="main", concretein=true)
function make_mlir_fn(mod, f, args, kwargs, name="main", concretein=true; toscalar=false)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Maybe we could turn name and concretein into kwargs?

Suggested change
function make_mlir_fn(mod, f, args, kwargs, name="main", concretein=true; toscalar=false)
function make_mlir_fn(mod, f, args, kwargs; name="main", concretein=true, toscalar=false)

src/overloads.jl Outdated
function elem_apply(f, args::VarArgs{Nargs})
primf = f.val

mod = MLIR.IR.module()
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
mod = MLIR.IR.module()
mod = MLIR.IR.mmodule()

src/overloads.jl Outdated Show resolved Hide resolved
@avik-pal
Copy link
Collaborator

avik-pal commented Jul 10, 2024

/home/wmoses/llvms/llvm16/install/bin/clang -I/usr/include/x86_64-linux-gnu/c++/11 -L/home/wmoses/llvms/llvm16/build/lib/x86_64-unknown-linux-gnu -stdlib=libc++ -v "$@"
need to be generalized? Same for the gcc and clang++ scripts

src/overloads.jl Outdated
OutShape = nothing
for (k, v) in seen_args
invmap[v] = k
OutShape size(k)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
OutShape size(k)
OutShape = size(k)

@wsmoses wsmoses merged commit cbeddbf into main Jul 14, 2024
6 of 13 checks passed
@wsmoses wsmoses deleted the batch branch July 14, 2024 01:17
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.

3 participants