Use `dlpack` for array interop #10

rejuvyesh · 2022-01-25T20:44:03Z

rejuvyesh · 2022-01-26T00:58:27Z

julia> judge(median(results), median(no_dlpack[1]))
2-element BenchmarkTools.BenchmarkGroup:
  tags: []
  "pytorchhub" => 4-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "bs=16" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+18.40% => regression)
                          "functorch" => TrialJudgement(+15.16% => regression)
                          "jl" => TrialJudgement(+585.51% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+10.65% => regression)
                          "functorch" => TrialJudgement(+244.91% => regression)
                          "jl" => TrialJudgement(+426.57% => regression)
          "bs=32" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+136.70% => regression)
                          "functorch" => TrialJudgement(+147.69% => regression)
                          "jl" => TrialJudgement(+841.88% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+122.57% => regression)
                          "functorch" => TrialJudgement(+132.32% => regression)
                          "jl" => TrialJudgement(+213.17% => regression)
          "bs=8" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+400.93% => regression)
                          "functorch" => TrialJudgement(+391.76% => regression)
                          "jl" => TrialJudgement(+699.24% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+427.55% => regression)
                          "functorch" => TrialJudgement(+403.57% => regression)
                          "jl" => TrialJudgement(+471.85% => regression)
          "bs=1" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+1582.03% => regression)
                          "functorch" => TrialJudgement(+1324.41% => regression)
                          "jl" => TrialJudgement(+1247.94% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(+7.27% => regression)
                          "functorch" => TrialJudgement(+7.67% => regression)
                          "jl" => TrialJudgement(-14.63% => improvement)
  "pytorchmlp" => 4-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "bs=16" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-12.30% => improvement)
                          "functorch" => TrialJudgement(-7.23% => improvement)
                          "jl" => TrialJudgement(+19.95% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-1.89% => invariant)
                          "functorch" => TrialJudgement(-2.81% => invariant)
                          "jl" => TrialJudgement(+13.12% => regression)
          "bs=32" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-17.17% => improvement)
                          "functorch" => TrialJudgement(-8.52% => improvement)
                          "jl" => TrialJudgement(+27.38% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-2.17% => invariant)
                          "functorch" => TrialJudgement(-4.80% => invariant)
                          "jl" => TrialJudgement(+19.82% => regression)
          "bs=8" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-15.03% => improvement)
                          "functorch" => TrialJudgement(-6.87% => improvement)
                          "jl" => TrialJudgement(+9.58% => regression)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-0.59% => invariant)
                          "functorch" => TrialJudgement(-2.66% => invariant)
                          "jl" => TrialJudgement(+6.15% => regression)
          "bs=1" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: []
                  "forward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-14.15% => improvement)
                          "functorch" => TrialJudgement(-7.62% => improvement)
                          "jl" => TrialJudgement(+2.07% => invariant)
                  "backward" => 3-element BenchmarkTools.BenchmarkGroup:
                          tags: []
                          "torch" => TrialJudgement(-4.48% => invariant)
                          "functorch" => TrialJudgement(-4.81% => invariant)
                          "jl" => TrialJudgement(-0.07% => invariant)

Might be doing this wrong but not beneficial on CPU? Need a better machine to evaluate.

rejuvyesh · 2022-01-28T23:19:13Z

Likely some issue with GC.@preserve. Need to figure out a MWE.

julia> for i in 1:10; TestEnv.activate() do; include("test/test_pytorch.jl"); end; end
Precompiling project...
  1 dependency successfully precompiled in 2 seconds (16 already precompiled)
Test Summary: | Pass  Total
dlpack        |    4      4
┌ Warning: `vendor()` is deprecated, use `BLAS.get_config()` and inspect the output instead
│   caller = npyinitialize() at numpy.jl:67
└ @ PyCall ~/.julia/packages/PyCall/L0fLP/src/numpy.jl:67
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:68
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.009436185 0.0996897 0.20193027; 0.38426346 -0.28908443 0.08106785], Float32[1.938506f-39 0.0 5.01105f-33; 0.0 4.5915f-41 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:68 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:68
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.096376784, -0.57178324], Float32[8.28208f-40, 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:68 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:73
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.009436185 0.0996897 0.20193027; 0.38426346 -0.28908443 0.08106785], Float32[1.938506f-39 0.0 5.01105f-33; 0.0 4.5915f-41 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:73 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:73
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.096376784, -0.57178324], Float32[8.28208f-40, 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:73 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:77
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.009436185 0.0996897 0.20193027; 0.38426346 -0.28908443 0.08106785], Float32[1.938506f-39 0.0 5.01105f-33; 0.0 4.5915f-41 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:77 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
linear: Test Failed at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:77
  Expression: isapprox(torchparams[i], linwrap.params[i], atol = 0.0001, rtol = 0.0001)
   Evaluated: isapprox(Float32[0.096376784, -0.57178324], Float32[8.28208f-40, 0.0]; atol = 0.0001, rtol = 0.0001)
Stacktrace:
 [1] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:445 [inlined]
 [2] macro expansion
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:77 [inlined]
 [3] macro expansion
   @ ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/Test/src/Test.jl:1283 [inlined]
 [4] top-level scope
   @ ~/.julia/dev/PyCallChainRules/test/test_pytorch.jl:59
Test Summary: | Pass  Fail  Total
linear        |   14     6     20
ERROR: LoadError: Some tests did not pass: 14 passed, 6 failed, 0 errored, 0 broken.
in expression starting at /home/jagupt/.julia/dev/PyCallChainRules/test/test_pytorch.jl:58

vs
All tests pass in:

julia> for i in 1:10; GC.enable(false); TestEnv.activate() do; include("test/test_pytorch.jl"); end; GC.enable(true) end

src/jax.jl

rejuvyesh · 2022-02-15T01:38:14Z

src/PyCallChainRules.jl

+
+maybecontiguous(x::AbstractArray) = Array(x)
+mayebecontiguous(x::StridedArray) = x
+function maybecontiguous(x::FillArrays.AbstractFill) 


This can't be the best way to handle FillArrays?

rejuvyesh mentioned this pull request Jan 26, 2022

Exposing dlpack interface for Julia Array/CuArray pabloferz/DLPack.jl#10

Closed

rejuvyesh added 3 commits January 27, 2022 22:51

Use dlpack for array interop

702ec12

more tests

22c5967

rebase but tests are flaky

65d8036

rejuvyesh force-pushed the jkg/dlpack branch from e2ceace to 65d8036 Compare January 27, 2022 22:58

rejuvyesh added 3 commits January 28, 2022 20:25

refactor to avoid repetition

890c517

missed import

a3b3995

small change to enable later use for gpu

42a1fc5

rejuvyesh mentioned this pull request Jan 29, 2022

Is GC.@preserve enough? pabloferz/DLPack.jl#14

Closed

rejuvyesh added 2 commits January 31, 2022 19:13

update for DLPack.jl refactor

ddf2b7e

update jax for new DLPack.jl

7888b40

pabloferz reviewed Jan 31, 2022

View reviewed changes

src/jax.jl Outdated Show resolved Hide resolved

pabloferz reviewed Jan 31, 2022

View reviewed changes

src/jax.jl Outdated Show resolved Hide resolved

rejuvyesh added 2 commits January 31, 2022 21:41

latest DLPack; all issues resolved

4d09900

minor cleanup

5e93561

pabloferz reviewed Jan 31, 2022

View reviewed changes

src/jax.jl Outdated Show resolved Hide resolved

rejuvyesh added 4 commits February 7, 2022 23:18

update for upcoming DLPack interface for sharing jlarrays to python

923a8b5

get jax cuda working

8e10397

start getting ready for GPUs

3de2ec2

fix for new version

fd81252

rejuvyesh mentioned this pull request Feb 15, 2022

Add share and wrap interfaces pabloferz/DLPack.jl#17

Merged

4 tasks

acknowledgement

afb9f5b

rejuvyesh commented Feb 15, 2022

View reviewed changes

rejuvyesh added 3 commits February 15, 2022 01:52

minor cleanup

ab108fc

use device

4b29709

fix for DLPack's share interface for PyCall

8dcff72

rejuvyesh marked this pull request as ready for review February 21, 2022 05:17

update version since we are DLPack based now

58c2b00

rejuvyesh added 10 commits February 21, 2022 05:19

add link

7158777

make CUDA optional

5eefee9

relax tolerance

3873560

relax CUDA versions

b2ef824

update version requirements

4c064a7

simply adapt

5ab994d

update readme for gpu

18c56f3

update readme

6493802

improve jax install instructions

fe00e4c

add basic kwargs support

c1223df

rejuvyesh merged commit 7b65d86 into main Feb 21, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use `dlpack` for array interop #10

Use `dlpack` for array interop #10

rejuvyesh commented Jan 25, 2022

rejuvyesh commented Jan 26, 2022

rejuvyesh commented Jan 28, 2022

rejuvyesh Feb 15, 2022

Use dlpack for array interop #10

Use dlpack for array interop #10

Conversation

rejuvyesh commented Jan 25, 2022

rejuvyesh commented Jan 26, 2022

rejuvyesh commented Jan 28, 2022

rejuvyesh Feb 15, 2022

Choose a reason for hiding this comment

Use `dlpack` for array interop #10

Use `dlpack` for array interop #10