JuliaDiff / JuliaDiff/ChainRules.jl
Pullback fails to inline with intervals
Open
Nobody has claimed this yet.
design
performance
- Dominant language
- Julia
- Stars
- 475
- Forks
- 98
- PR merge metrics
- No merged PRs in 30d
Description
The following MWE does not inline when using intervals from IntervalArithmetic.jl:
function fff(x, y)
z, z_pullback = rrule(*, x, y)
z̄ = one(x)
_, r1, r2 = z_pullback(z̄)
x̄ = unthunk(r1)
ȳ = unthunk(r2)
return (x̄, ȳ)
end
using IntervalArithmetic
julia> @code_native f(1..1, 2..2)
.section __TEXT,__text,regular,pure_instructions
; ┌ @ REPL[23]:1 within `f'
pushq %rbx
subq $160, %rsp
movq %rdi, %rbx
; │ @ REPL[23]:2 within `f'
movabsq $rrule, %rax
leaq 112(%rsp), %rdi
callq *%rax
movabsq $5516288192, %rax ## imm = 0x148CBE0C0
; │ @ REPL[23]:5 within `f'
; │┌ @ fastmath_able.jl:188 within `times_pullback'
vmovaps (%rax), %xmm0
vmovups %xmm0, 80(%rsp)
vmovups 128(%rsp), %xmm1
vmovups 144(%rsp), %xmm2
vmovups %xmm2, 96(%rsp)
vmovups %xmm0, 48(%rsp)
vmovups %xmm1, 64(%rsp)
; │└
; │ @ REPL[23]:7 within `f'
; │┌ @ thunks.jl:99 within `unthunk'
; ││┌ @ thunks.jl:98 within `Thunk'
movabsq $"#461", %rax
leaq 16(%rsp), %rdi
leaq 80(%rsp), %rsi
callq *%rax
; │└└
; │ @ REPL[23]:8 within `f'
; │┌ @ thunks.jl:99 within `unthunk'
; ││┌ @ thunks.jl:98 within `Thunk'
movabsq $"#462", %rax
movq %rsp, %rdi
leaq 48(%rsp), %rsi
callq *%rax
; │└└
; │ @ REPL[23]:10 within `f'
vmovups (%rsp), %xmm0
vmovups %xmm0, 32(%rsp)
vmovups 16(%rsp), %ymm0
vmovups %ymm0, (%rbx)
movq %rbx, %rax
addq $160, %rsp
popq %rbx
vzeroupper
retq
nopw %cs:(%rax,%rax)
nopl (%rax)
; └
Cf. the beautiful code when using floats:
julia> @code_native f(1.0, 2.0)
.section __TEXT,__text,regular,pure_instructions
; ┌ @ REPL[23]:1 within `f'
movq %rdi, %rax
; │ @ REPL[23]:10 within `f'
vmovsd %xmm1, (%rdi)
vmovsd %xmm0, 8(%rdi)
retq
nopl (%rax)
; └
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the MWE using IntervalArithmetic.jl, the fff function, rrule(*, x, y), and @code_native. Compare the interval output with the float output and trace why the pullback and unthunk calls remain uninlined. Done means the interval case produces appropriately inlined native code without regressing the float case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100