Add tessera option to @trace
#1986
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
``` @trace tessera mysin(x)=sin(x) ``` This will add an attribute with name `tessera_name` and value a stringattr with the function name (`"mysin"`) to the generated function for mysin. (note that you need `@code_hlo optimize=false` to see this attribute, otherwise it is inlined.)
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
lib/ReactantCore/src/ReactantCore.jl
Outdated
Comment on lines
170
to
171
| tessera = true | ||
| args = args[2:end] |
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
| tessera = true | |
| args = args[2:end] | |
| tessera = true | |
| args = args[2:end] |
lib/ReactantCore/src/ReactantCore.jl
Outdated
| function trace_function_definition(mod, expr; tessera=false) | ||
| internal_fn = MacroTools.splitdef(expr) | ||
| orig_fname = internal_fn[:name] | ||
|
|
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
lib/ReactantCore/src/ReactantCore.jl
Outdated
|
|
||
| if isempty(new_fn[:kwargs]) | ||
| traced_call_expr = :($(traced_call)($(fname), $(argnames...))) | ||
| traced_call_expr = :($(traced_call)($(fname), $(argnames...); tessera_name=$(String(tessera_name)))) |
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
| traced_call_expr = :($(traced_call)($(fname), $(argnames...); tessera_name=$(String(tessera_name)))) | |
| traced_call_expr = :($(traced_call)( | |
| $(fname), $(argnames...); tessera_name=$(String(tessera_name)) | |
| )) |
lib/ReactantCore/src/ReactantCore.jl
Outdated
Comment on lines
283
to
284
| traced_call_expr = | ||
| :($(traced_call)(Core.kwcall, (; $(kws...)), $(fname), $(argnames...))) | ||
| :($(traced_call)(Core.kwcall, (; $(kws...)), $(fname), $(argnames...); tessera_name=$(String(tessera_name)))) |
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
| traced_call_expr = | |
| :($(traced_call)(Core.kwcall, (; $(kws...)), $(fname), $(argnames...))) | |
| :($(traced_call)(Core.kwcall, (; $(kws...)), $(fname), $(argnames...); tessera_name=$(String(tessera_name)))) | |
| traced_call_expr = :($(traced_call)( | |
| Core.kwcall, | |
| (; $(kws...)), | |
| $(fname), | |
| $(argnames...); | |
| tessera_name=$(String(tessera_name)), | |
| )) |
src/Ops.jl
Outdated
| end | ||
|
|
||
| @noinline function call(f, args...; location=mlir_stacktrace("call", @__FILE__, @__LINE__)) | ||
| @noinline function call(f, args...; location=mlir_stacktrace("call", @__FILE__, @__LINE__), tessera_name=nothing) |
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
| @noinline function call(f, args...; location=mlir_stacktrace("call", @__FILE__, @__LINE__), tessera_name=nothing) | |
| @noinline function call( | |
| f, args...; location=mlir_stacktrace("call", @__FILE__, @__LINE__), tessera_name=nothing | |
| ) |
src/Ops.jl
Outdated
| resargprefix, | ||
| ) | ||
| if !isnothing(tessera_name) | ||
| MLIR.IR.attr!(temp.f, "tessera_name", MLIR.IR.Attribute(tessera_name)) |
Contributor
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
Suggested change
| MLIR.IR.attr!(temp.f, "tessera_name", MLIR.IR.Attribute(tessera_name)) | |
| MLIR.IR.attr!(temp.f, "tessera_name", MLIR.IR.Attribute(tessera_name)) |
jumerckx
force-pushed
the
tessera_defs
branch
from
3efbed4 to
d862e5a
Compare
jumerckx
force-pushed
the
tessera_defs
branch
from
d862e5a to
6804c97
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This will add an attribute with name
tessera_nameand value a stringattr with the function name ("reciprocal") to the generated MLIR function. (note that you need@code_hlo optimize=falseto see this attribute, otherwise it is inlined.)generates code containing:
(note the
tessera_nameattribute)It might be nicer to have a dedicated macro, i.e.:
I just did it this way because it was quicker to reuse some stuff.