From 9f94ffc816f54355fce98e1b2b718e7703933403 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 2 Jul 2024 23:26:27 -0700 Subject: [PATCH 1/9] chore: update compat entries --- Project.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 59be99e..e4f5472 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SimpleChains" uuid = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5" authors = ["Chris Elrod and contributors"] -version = "0.4.6" +version = "0.4.7" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" @@ -40,10 +40,10 @@ ManualMemory = "0.1.8" Polyester = "0.4, 0.5, 0.6, 0.7" SIMDTypes = "0.1" SLEEFPirates = "0.6" -Static = "0.8.4" +Static = "0.8.4, 1" StaticArrayInterface = "1" StaticArrays = "1" -StrideArraysCore = "0.4.7" +StrideArraysCore = "0.4.7, 0.5" UnPack = "1" VectorizationBase = "0.21.40" VectorizedRNG = "0.2.13" From 4277b26022d949ac3f906475b32a6d8ae6eb9176 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jul 2024 20:39:46 -0700 Subject: [PATCH 2/9] fix: update +(Ptr, StaticInt) to __add --- .github/workflows/CI.yml | 1 - Project.toml | 6 +++--- src/chain_rules.jl | 4 ++-- src/dropout.jl | 2 +- src/simple_chain.jl | 4 ++-- src/utils.jl | 4 ++++ test/runtests.jl | 16 +++------------- 7 files changed, 15 insertions(+), 22 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c533927..c5f1433 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -81,7 +81,6 @@ jobs: os: - ubuntu-latest version: - - '1.6' - '1' steps: - uses: actions/checkout@c85c95e3d7251135ab7dc9ce3241c5835cc595a9 diff --git a/Project.toml b/Project.toml index e4f5472..5008c10 100644 --- a/Project.toml +++ b/Project.toml @@ -40,14 +40,14 @@ ManualMemory = "0.1.8" Polyester = "0.4, 0.5, 0.6, 0.7" SIMDTypes = "0.1" SLEEFPirates = "0.6" -Static = "0.8.4, 1" +Static = "1" StaticArrayInterface = "1" StaticArrays = "1" -StrideArraysCore = "0.4.7, 0.5" +StrideArraysCore = "0.5" UnPack = "1" VectorizationBase = "0.21.40" VectorizedRNG = "0.2.13" -julia = "1.6" +julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" diff --git a/src/chain_rules.jl b/src/chain_rules.jl index 2a124db..72f1b20 100644 --- a/src/chain_rules.jl +++ b/src/chain_rules.jl @@ -162,9 +162,9 @@ function valgrad_noloss( @inbounds @simd ivdep for i in eachindex(parg) parg2[i] = parg[i] end - pm += aoff + pm = __add(pm, aoff) g = PtrArray(Ptr{T}(pm), (glen,)) - pm += goff + pm = __add(pm, goff) l, pbl = chain_valgrad_pullback!(pointer(g), parg2, layers, pointer(params), pm) end diff --git a/src/dropout.jl b/src/dropout.jl index 4cefd08..2ac46a0 100644 --- a/src/dropout.jl +++ b/src/dropout.jl @@ -120,7 +120,7 @@ function valgrad_layer!( VectorizedRNG.storestate!(rng, state) end # GC preserve - pg, x, p, align(pu + ((static(7) + N) >>> static(3))) + pg, x, p, align(__add(pu, ((static(7) + N) >>> static(3)))) end function pullback_arg!( diff --git a/src/simple_chain.jl b/src/simple_chain.jl index fdd43fc..4cda5f3 100644 --- a/src/simple_chain.jl +++ b/src/simple_chain.jl @@ -820,7 +820,7 @@ function valgrad_core( ) where {T} @unpack layers = c g = PtrArray(Ptr{T}(pu), (glen,)) - l = unsafe_valgrad!(c, pu + align(glen * static_sizeof(T)), g, params, arg) + l = unsafe_valgrad!(c, __add(pu, align(glen * static_sizeof(T))), g, params, arg) Base.FastMath.add_fast( l, apply_penalty!(g, getpenalty(c), params, static_size(arg)) @@ -838,7 +838,7 @@ function valgrad_core_sarray( l = Base.FastMath.add_fast( unsafe_valgrad!( c, - pu + align(static(L) * static_sizeof(T)), + __add(pu, align(static(L) * static_sizeof(T))), g, params, arg diff --git a/src/utils.jl b/src/utils.jl index d56240d..e754956 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -336,3 +336,7 @@ function _add_memory(t::Tuple, p) (A, B...) end _add_memory(::Nothing, p) = nothing + +__add(x, y) = x + y +__add(x::Ptr, ::StaticInt{N}) where {N} = x + N +__add(::StaticInt{N}, y::Ptr) where {N} = y + N diff --git a/test/runtests.jl b/test/runtests.jl index 67b027f..f8a7574 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,12 +1,6 @@ using SimpleChains using Test, Aqua, ForwardDiff, Zygote, ChainRules, Random -@static if VERSION >= v"1.9" - using JET: @test_opt -else - macro test_opt(ex) - nothing - end -end +using JET: @test_opt countallocations!(g, sc, x, p) = @allocated valgrad!(g, sc, x, p) dual(x::T) where {T} = ForwardDiff.Dual(x, 4randn(T), 4randn(T), 4randn(T)) @@ -84,12 +78,8 @@ InteractiveUtils.versioninfo(; verbose = true) SquaredLoss""" @test sprint((io, t) -> show(io, t), sc) == print_str0 - if VERSION >= v"1.6" - @test sprint((io, t) -> show(io, t), scflp) == print_str1 - else - # typename doesn't work on 1.5 - @test_broken sprint((io, t) -> show(io, t), scflp) == print_str1 - end + @test sprint((io, t) -> show(io, t), scflp) == print_str1 + p = SimpleChains.init_params(scflp, T; rng = Random.default_rng()) g = similar(p) let sc = SimpleChains.remove_loss(sc) From 074db053dbf97bd4b618646577f2ddf5dc696617 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jul 2024 21:14:37 -0700 Subject: [PATCH 3/9] fix: special case adjoint for staticvector --- src/SimpleChains.jl | 1 + src/dense.jl | 8 ++++---- src/utils.jl | 9 ++++++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/SimpleChains.jl b/src/SimpleChains.jl index 4ac9f46..461d053 100644 --- a/src/SimpleChains.jl +++ b/src/SimpleChains.jl @@ -42,6 +42,7 @@ import ChainRulesCore import ForwardDiff import LoopVectorization import StaticArrays +using StaticArrays: SVector, SMatrix using Random: AbstractRNG using LoopVectorization: matmul_params, @turbo diff --git a/src/dense.jl b/src/dense.jl index c9bde99..273d433 100644 --- a/src/dense.jl +++ b/src/dense.jl @@ -1004,7 +1004,7 @@ end function dense_param_update!(::TurboDense{true}, Ā, C̄, B) Kp1 = static_size(Ā, StaticInt(2)) K = Kp1 - StaticInt(1) - dense!(identity, nothing, view(Ā, :, static(1):K), C̄, B', False()) + dense!(identity, nothing, view(Ā, :, static(1):K), C̄, __adjoint(B), False()) @turbo for m ∈ axes(Ā, 1) s = zero(eltype(Ā)) for n ∈ axes(C̄, 2) @@ -1014,12 +1014,12 @@ function dense_param_update!(::TurboDense{true}, Ā, C̄, B) end end function dense_param_update!(::TurboDense{false}, Ā, C̄, B) - dense!(identity, nothing, Ā, C̄, B', False()) + dense!(identity, nothing, Ā, C̄, __adjoint(B), False()) end function dense_param_update!(::TurboDense{true}, Ā, C̄, B, inds) Kp1 = static_size(Ā, StaticInt(2)) K = Kp1 - StaticInt(1) - denserev!(identity, nothing, view(Ā, :, static(1):K), C̄, B', inds, False()) + denserev!(identity, nothing, view(Ā, :, static(1):K), C̄, __adjoint(B), inds, False()) @turbo for m ∈ axes(Ā, 1) s = zero(eltype(Ā)) for n ∈ axes(C̄, 2) @@ -1029,7 +1029,7 @@ function dense_param_update!(::TurboDense{true}, Ā, C̄, B, inds) end end function dense_param_update!(::TurboDense{false}, Ā, C̄, B, inds) - dense!(identity, nothing, Ā, C̄, B', inds, False()) + dense!(identity, nothing, Ā, C̄, __adjoint(B), inds, False()) end @inline function dense!( diff --git a/src/utils.jl b/src/utils.jl index e754956..06fc867 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -337,6 +337,9 @@ function _add_memory(t::Tuple, p) end _add_memory(::Nothing, p) = nothing -__add(x, y) = x + y -__add(x::Ptr, ::StaticInt{N}) where {N} = x + N -__add(::StaticInt{N}, y::Ptr) where {N} = y + N +@inline __add(x, y) = x + y +@inline __add(x::Ptr, ::StaticInt{N}) where {N} = x + N +@inline __add(::StaticInt{N}, y::Ptr) where {N} = y + N + +@inline __adjoint(x) = x' +@inline __adjoint(x::SVector{N, <:Real}) where {N} = SMatrix{1, N, eltype(x)}(x.data) From c08a84b8c271fe81ba1547afd4d8f8f9837411c5 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jul 2024 21:25:59 -0700 Subject: [PATCH 4/9] fix: update Aqua + doc link --- docs/src/examples/mnist.md | 2 +- test/runtests.jl | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/src/examples/mnist.md b/docs/src/examples/mnist.md index 44e6938..6e888a2 100644 --- a/docs/src/examples/mnist.md +++ b/docs/src/examples/mnist.md @@ -43,7 +43,7 @@ We define the inputs as being statically sized `(28,28,1)` images. Specifying the input sizes allows these to be checked. Making them static, which we can do either in our simple chain, or by adding static sizing to the images themselves using a package like [StrideArrays.jl](https://github.com/JuliaSIMD/StrideArrays.jl) -or [HybridArrays.jl](git@github.com:JuliaArrays/HybridArrays.jl.git). These packages are recommended +or [HybridArrays.jl](https://github.com/JuliaArrays/HybridArrays.jl). These packages are recommended for allowing you to mix dynamic and static sizes; the batch size should probably be left dynamic, as you're unlikely to want to specialize code generation on this, given that it is likely to vary, increasing compile times while being unlikely to diff --git a/test/runtests.jl b/test/runtests.jl index f8a7574..f6510d0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -564,5 +564,4 @@ end Aqua.test_all( SimpleChains; ambiguities = false, - project_toml_formatting = false ) From 70f28fc8b0026beae209f395c54251d643e6af14 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jul 2024 21:28:19 -0700 Subject: [PATCH 5/9] chore: remove unused extra deps --- Project.toml | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index 5008c10..ebb59cf 100644 --- a/Project.toml +++ b/Project.toml @@ -38,6 +38,7 @@ LayoutPointers = "0.1.3" LoopVectorization = "0.12.104" ManualMemory = "0.1.8" Polyester = "0.4, 0.5, 0.6, 0.7" +Random = "1.10" SIMDTypes = "0.1" SLEEFPirates = "0.6" Static = "1" @@ -48,12 +49,3 @@ UnPack = "1" VectorizationBase = "0.21.40" VectorizedRNG = "0.2.13" julia = "1.10" - -[extras] -Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[targets] -test = ["Aqua", "ForwardDiff", "Zygote", "Test"] From 2c337ead6573301a8238a97e8bb71fa43f529c35 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jul 2024 08:32:08 -0700 Subject: [PATCH 6/9] chore: remove forced inlining --- src/utils.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index 06fc867..b155b07 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -337,9 +337,9 @@ function _add_memory(t::Tuple, p) end _add_memory(::Nothing, p) = nothing -@inline __add(x, y) = x + y -@inline __add(x::Ptr, ::StaticInt{N}) where {N} = x + N -@inline __add(::StaticInt{N}, y::Ptr) where {N} = y + N +__add(x, y) = x + y +__add(x::Ptr, ::StaticInt{N}) where {N} = x + N +__add(::StaticInt{N}, y::Ptr) where {N} = y + N -@inline __adjoint(x) = x' -@inline __adjoint(x::SVector{N, <:Real}) where {N} = SMatrix{1, N, eltype(x)}(x.data) +__adjoint(x) = x' +__adjoint(x::SVector{N, <:Real}) where {N} = SMatrix{1, N, eltype(x)}(x.data) From 10886ba0efb58107dd7817a4c798263edc708546 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jul 2024 08:32:31 -0700 Subject: [PATCH 7/9] ci: coverage for julia:1 --- .github/workflows/CI.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c5f1433..d9951dc 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -62,7 +62,7 @@ jobs: - uses: julia-actions/julia-buildpkg@72ddd0fcdc760e6d2ef7ef2bbb11ae5032b671a8 - uses: julia-actions/julia-runtest@79a7e100883947123f8263c5f06e6c0ea3eb972f with: - coverage: true + coverage: false env: JULIA_NUM_THREADS: ${{ matrix.threads }} JULIA_CPU_THREADS: ${{ matrix.threads }} @@ -101,7 +101,7 @@ jobs: - uses: julia-actions/julia-buildpkg@72ddd0fcdc760e6d2ef7ef2bbb11ae5032b671a8 - uses: julia-actions/julia-runtest@79a7e100883947123f8263c5f06e6c0ea3eb972f with: - coverage: false + coverage: true env: JULIA_NUM_THREADS: 2 docs: From b087b76e5d57c914832d1cc3a321712351c87c15 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jul 2024 08:39:52 -0700 Subject: [PATCH 8/9] chore: check older versions for 1.6 --- .github/workflows/CI.yml | 7 ++++--- Project.toml | 6 +++--- test/runtests.jl | 8 +++++++- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index d9951dc..ad5ff08 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -39,6 +39,7 @@ jobs: os: - ubuntu-latest version: + - '1' - 'nightly' # coverage fast on nightly threads: - '3' @@ -62,7 +63,7 @@ jobs: - uses: julia-actions/julia-buildpkg@72ddd0fcdc760e6d2ef7ef2bbb11ae5032b671a8 - uses: julia-actions/julia-runtest@79a7e100883947123f8263c5f06e6c0ea3eb972f with: - coverage: false + coverage: true env: JULIA_NUM_THREADS: ${{ matrix.threads }} JULIA_CPU_THREADS: ${{ matrix.threads }} @@ -81,7 +82,7 @@ jobs: os: - ubuntu-latest version: - - '1' + - '1.6' steps: - uses: actions/checkout@c85c95e3d7251135ab7dc9ce3241c5835cc595a9 - uses: julia-actions/setup-julia@f40c4b69330df1d22e7590c12e76dc2f9c66e0bc @@ -101,7 +102,7 @@ jobs: - uses: julia-actions/julia-buildpkg@72ddd0fcdc760e6d2ef7ef2bbb11ae5032b671a8 - uses: julia-actions/julia-runtest@79a7e100883947123f8263c5f06e6c0ea3eb972f with: - coverage: true + coverage: false env: JULIA_NUM_THREADS: 2 docs: diff --git a/Project.toml b/Project.toml index ebb59cf..705f4f0 100644 --- a/Project.toml +++ b/Project.toml @@ -38,14 +38,14 @@ LayoutPointers = "0.1.3" LoopVectorization = "0.12.104" ManualMemory = "0.1.8" Polyester = "0.4, 0.5, 0.6, 0.7" -Random = "1.10" +Random = "<0.0.1, 1" SIMDTypes = "0.1" SLEEFPirates = "0.6" -Static = "1" +Static = "0.8.4, 1" StaticArrayInterface = "1" StaticArrays = "1" StrideArraysCore = "0.5" UnPack = "1" VectorizationBase = "0.21.40" VectorizedRNG = "0.2.13" -julia = "1.10" +julia = "1.6" diff --git a/test/runtests.jl b/test/runtests.jl index f6510d0..8edc9fb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,12 @@ using SimpleChains using Test, Aqua, ForwardDiff, Zygote, ChainRules, Random -using JET: @test_opt +@static if VERSION ≥ v"1.9" + using JET: @test_opt +else + macro test_opt(ex) + nothing + end +end countallocations!(g, sc, x, p) = @allocated valgrad!(g, sc, x, p) dual(x::T) where {T} = ForwardDiff.Dual(x, 4randn(T), 4randn(T), 4randn(T)) From 7e02c02284c5fdaa6f39d4194013e80f31bf69c8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jul 2024 00:08:25 -0700 Subject: [PATCH 9/9] chore: suggestions from code-review --- .github/workflows/CI.yml | 1 - Project.toml | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index ad5ff08..05a9754 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -40,7 +40,6 @@ jobs: - ubuntu-latest version: - '1' - - 'nightly' # coverage fast on nightly threads: - '3' - '4' diff --git a/Project.toml b/Project.toml index 705f4f0..48da169 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,7 @@ SLEEFPirates = "0.6" Static = "0.8.4, 1" StaticArrayInterface = "1" StaticArrays = "1" -StrideArraysCore = "0.5" +StrideArraysCore = "0.4.7, 0.5" UnPack = "1" VectorizationBase = "0.21.40" VectorizedRNG = "0.2.13"