Improve _apply(apply_type, ::Tuple, ::SimpleVector) #29922
Conversation
|
Hmm, passes tests, but seems to be causing some miscompiles elsewhere. Looking into it. |
Keno
force-pushed
the
kf/_apply_apply_type
branch
from
abb60f4 to
2e07f11
Compare
|
Never mind, I did a stupid. Fixed. |
|
(For those following along at home, this fixes compiling the code generated by Zygote for the |
base/compiler/ssair/inlining.jl
Outdated
| new_call = Expr(:call, Core.getfield, def, j) | ||
| new_argexpr = insert_node!(ir, idx, def_atype, new_call) | ||
| if isa(def_atype, Const) && is_inlineable_constant(def_atype.val) | ||
| new_argexpr = insert_node!(ir, idx, def_atype, quoted(def_atype.val)) |
There was a problem hiding this comment.
Const shouldn't need a node at all, right, just use the value as the argexpr?
There was a problem hiding this comment.
yes, I suspect that's true. I'll try it.
test/compiler/compiler.jl
Outdated
| end | ||
| let ci = code_typed(foo_apply_apply_type_svec, Tuple{})[1].first | ||
| @test length(ci.code) == 1 | ||
| end |
There was a problem hiding this comment.
can we make a better test than this? this one seems brittle under various compiler flags (notable, the code-coverage bot, although code_typed itself likes to mislead about this particular form of result also)—although perhaps this is really the best we can do, since we need to check for a particular optimization result?
perhaps we could at least improve the test to be more explicit: ci.code == Any[ Expr(:return, Const(NTuple{3, Float32})) ] (e.g. the expression that code_typed may or may not decide to replace the optimization result with)
|
Maybe not strictly relevant to this PR, but also ref https://github.com/jrevels/Cassette.jl/blob/cfa248a99209d17e8d1997dd887a698c4bc079b9/src/context.jl#L444 (in case you run into needing a similar specialization hack for your TPU purposes) |
Keno
force-pushed
the
kf/_apply_apply_type
branch
from
2e07f11 to
b2bff7f
Compare
This is another one of those patterns that doesn't generally come up except in Cassette/Zygote code. There's two related changes here: 1. Expand ininling's _apply rewrite to also handle constant svecs (under the restriction that they must be no longer than the splatting cutoff and the elements must be individually eligible for inlining into the IR) 2. Move the _apply rewrite before the special case inliner for builtins such that _apply(apply_type, ...) gets eliminated early.
Keno
force-pushed
the
kf/_apply_apply_type
branch
from
b2bff7f to
92ac4ac
Compare
|
Somehow I though this had gotten merged and then got a rude awakening when this caused me bad performance when evaluating #31253. Rebased and addressed review comments. |
This is another one of those patterns that doesn't generally
come up except in Cassette/Zygote code. There's two related
changes here:
svecs (under the restriction that they must be no longer than
the splatting cutoff and the elements must be individually eligible
for inlining into the IR)
such that _apply(apply_type, ...) gets eliminated early.