Skip to content

Commit

Permalink
Float32 Testing (#414)
Browse files Browse the repository at this point in the history 
* Extend testing utils for f32

* Test basic maths at f32

* Test BLAS at f32

* Bump patch

* Builtins

* fastmath

* Simplify BLAS testing

* flat_product

* lapack testing

* Tidy up blas testing

* Helper functions to simplify testing

* Tidy up blas and lapack rules

* matrix exponential on f32

* Fix numerical issue

* Formatting

* Fix up MatrixBeta test

* SpecialFunctions on F32

* Fix array integration testing

* Test LogExpFunctions on F32

* Enable Lux correctness testing

* Fix LAPACK rules

* Try a different method

* Formatting

* Tailor perf for 1.10

* Improve trsm stability

* Simplify BLAS and LAPACK testing

* Use StableRNG everywhere

* Formatting

* Style + formatting
  • Loading branch information
@willtebbutt
willtebbutt authored Dec 11, 2024
1 parent e581487 commit 102f900
Show file tree
Hide file tree
Showing 14 changed files with 787 additions and 911 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "Mooncake"
uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
authors = ["Will Tebbutt, Hong Ge, and contributors"]
version = "0.4.59"
version = "0.4.60"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand Down
218 changes: 95 additions & 123 deletions src/rrules/blas.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -720,7 +720,6 @@ function blas_vectors(rng::AbstractRNG, P::Type{<:BlasRealFloat}, p::Int)
xs = Any[
randn(rng, P, p),
view(randn(rng, P, p + 5), 3:(p + 2)),
view(randn(rng, P, 3p), 1:2:(2p)),
view(randn(rng, P, 3p, 3), 1:2:(2p), 2),
]
@assert all(x -> length(x) == p, xs)
Expand All @@ -732,72 +731,52 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:blas})
t_flags = ['N', 'T', 'C']
alphas = [1.0, -0.25]
betas = [0.0, 0.33]
Ps = [Float64, Float32]
rng = rng_ctor(123456)

test_cases = vcat(

# gemv!
vec(
reduce(
vcat,
map(product(t_flags, [1, 3], [1, 2])) do (tA, M, N)
t = tA == 'N'
As = blas_matrices(rng, Float64, t ? M : N, t ? N : M)
xs = blas_vectors(rng, Float64, N)
ys = blas_vectors(rng, Float64, M)
flags = (false, :stability, (lb=1e-3, ub=10.0))
return map(product(As, xs, ys)) do (A, x, y)
return (flags..., BLAS.gemv!, tA, randn(), A, x, randn(), y)
end
end,
),
),
map_prod(t_flags, [1, 3], [1, 2], Ps) do (tA, M, N, P)
As = blas_matrices(rng, P, tA == 'N' ? M : N, tA == 'N' ? N : M)
xs = blas_vectors(rng, P, N)
ys = blas_vectors(rng, P, M)
flags = (false, :stability, (lb=1e-3, ub=10.0))
return map(As, xs, ys) do A, x, y
return (flags..., BLAS.gemv!, tA, randn(P), A, x, randn(P), y)
end
end...,

# symv!
vec(
reduce(
vcat,
map(product(['L', 'U'], alphas, betas)) do (uplo, α, β)
As = blas_matrices(rng, Float64, 5, 5)
ys = blas_vectors(rng, Float64, 5)
xs = blas_vectors(rng, Float64, 5)
return map(product(As, xs, ys)) do (A, x, y)
(false, :stability, nothing, BLAS.symv!, uplo, α, A, x, β, y)
end
end,
),
),
map_prod(['L', 'U'], alphas, betas, Ps) do (uplo, α, β, P)
As = blas_matrices(rng, P, 5, 5)
ys = blas_vectors(rng, P, 5)
xs = blas_vectors(rng, P, 5)
return map(As, xs, ys) do A, x, y
(false, :stability, nothing, BLAS.symv!, uplo, P(α), A, x, P(β), y)
end
end...,

# gemm!
vec(
reduce(
vcat,
map(product(t_flags, t_flags, alphas, betas)) do (tA, tB, a, b)
As = blas_matrices(rng, Float64, tA == 'N' ? 3 : 4, tA == 'N' ? 4 : 3)
Bs = blas_matrices(rng, Float64, tB == 'N' ? 4 : 5, tB == 'N' ? 5 : 4)
Cs = blas_matrices(rng, Float64, 3, 5)
return map(product(As, Bs, Cs)) do (A, B, C)
(false, :none, nothing, BLAS.gemm!, tA, tB, a, A, B, b, C)
end
end,
),
),
map_prod(t_flags, t_flags, alphas, betas, Ps) do (tA, tB, a, b, P)
As = blas_matrices(rng, P, tA == 'N' ? 3 : 4, tA == 'N' ? 4 : 3)
Bs = blas_matrices(rng, P, tB == 'N' ? 4 : 5, tB == 'N' ? 5 : 4)
Cs = blas_matrices(rng, P, 3, 5)
return map(As, Bs, Cs) do A, B, C
(false, :none, nothing, BLAS.gemm!, tA, tB, P(a), A, B, P(b), C)
end
end...,

# symm!
vec(
reduce(
vcat,
map(product(['L', 'R'], ['L', 'U'], alphas, betas)) do (side, uplo, α, β)
nA = side == 'L' ? 5 : 7
As = blas_matrices(rng, Float64, nA, nA)
Bs = blas_matrices(rng, Float64, 5, 7)
Cs = blas_matrices(rng, Float64, 5, 7)
return map(product(As, Bs, Cs)) do (A, B, C)
(false, :stability, nothing, BLAS.symm!, side, uplo, α, A, B, β, C)
end
end,
),
),
map_prod(['L', 'R'], ['L', 'U'], alphas, betas, Ps) do (side, ul, α, β, P)
nA = side == 'L' ? 5 : 7
As = blas_matrices(rng, P, nA, nA)
Bs = blas_matrices(rng, P, 5, 7)
Cs = blas_matrices(rng, P, 5, 7)
return map(As, Bs, Cs) do A, B, C
(false, :stability, nothing, BLAS.symm!, side, ul, P(α), A, B, P(β), C)
end
end...,
)

memory = Any[]
Expand All @@ -807,6 +786,10 @@ end
function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:blas})
t_flags = ['N', 'T', 'C']
aliased_gemm! = (tA, tB, a, b, A, C) -> BLAS.gemm!(tA, tB, a, A, A, b, C)
Ps = [Float32, Float64]
uplos = ['L', 'U']
dAs = ['N', 'U']
rng = rng_ctor(123)

test_cases = vcat(

Expand All @@ -820,91 +803,80 @@ function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:blas})
# BLAS LEVEL 1
#

Any[
(false, :none, nothing, BLAS.dot, 3, randn(5), 1, randn(4), 1),
(false, :none, nothing, BLAS.dot, 3, randn(6), 2, randn(4), 1),
(false, :none, nothing, BLAS.dot, 3, randn(6), 1, randn(9), 3),
(false, :none, nothing, BLAS.dot, 3, randn(12), 3, randn(9), 2),
(false, :none, nothing, BLAS.scal!, 10, 2.4, randn(30), 2),
],
map(Ps) do P
Any[
(false, :none, nothing, BLAS.dot, 3, randn(P, 5), 1, randn(P, 4), 1),
(false, :none, nothing, BLAS.dot, 3, randn(P, 6), 2, randn(P, 4), 1),
(false, :none, nothing, BLAS.dot, 3, randn(P, 6), 1, randn(P, 9), 3),
(false, :none, nothing, BLAS.dot, 3, randn(P, 12), 3, randn(P, 9), 2),
(false, :none, nothing, BLAS.scal!, 10, P(2.4), randn(P, 30), 2),
]
end...,

#
# BLAS LEVEL 2
#

# trmv!
vec(
reduce(
vcat,
map(product(['L', 'U'], t_flags, ['N', 'U'], [1, 3])) do (ul, tA, dA, N)
As = [randn(N, N), view(randn(15, 15), 3:(N + 2), 4:(N + 3))]
bs = [randn(N), view(randn(14), 4:(N + 3))]
return map(product(As, bs)) do (A, b)
(false, :none, nothing, BLAS.trmv!, ul, tA, dA, A, b)
end
end,
),
),
map_prod(uplos, t_flags, dAs, [1, 3], Ps) do (ul, tA, dA, N, P)
As = blas_matrices(rng, P, N, N)
bs = blas_vectors(rng, P, N)
return map(As, bs) do A, b
(false, :none, nothing, BLAS.trmv!, ul, tA, dA, A, b)
end
end...,

#
# BLAS LEVEL 3
#

# aliased gemm!
vec(
map(product(t_flags, t_flags)) do (tA, tB)
A = randn(5, 5)
B = randn(5, 5)
(false, :none, nothing, aliased_gemm!, tA, tB, randn(), randn(), A, B)
end,
),
map_prod(t_flags, t_flags, Ps) do (tA, tB, P)
As = blas_matrices(rng, P, 5, 5)
Bs = blas_matrices(rng, P, 5, 5)
return map_prod(As, Bs) do (A, B)
(false, :none, nothing, aliased_gemm!, tA, tB, randn(P), randn(P), A, B)
end
end...,

# syrk!
vec(
map(product(['U', 'L'], t_flags)) do (uplo, t)
A = t == 'N' ? randn(3, 4) : randn(4, 3)
C = randn(3, 3)
return (false, :none, nothing, BLAS.syrk!, uplo, t, randn(), A, randn(), C)
end,
),
map_prod(uplos, t_flags, Ps) do (uplo, t, P)
As = blas_matrices(rng, P, t == 'N' ? 3 : 4, t == 'N' ? 4 : 3)
C = randn(P, 3, 3)
flags = (false, :none, nothing)
return map(As) do A
return (flags..., BLAS.syrk!, uplo, t, randn(P), A, randn(P), C)
end
end...,

# trmm!
vec(
reduce(
vcat,
map(
product(['L', 'R'], ['U', 'L'], t_flags, ['N', 'U'], [1, 3], [1, 2])
) do (side, ul, tA, dA, M, N)
t = tA == 'N'
R = side == 'L' ? M : N
As = [randn(R, R), view(randn(15, 15), 3:(R + 2), 4:(R + 3))]
Bs = [randn(M, N), view(randn(15, 15), 2:(M + 1), 5:(N + 4))]
flags = (false, :none, nothing)
return map(product(As, Bs)) do (A, B)
(flags..., BLAS.trmm!, side, ul, tA, dA, randn(), A, B)
end
end,
),
),
map_prod(
['L', 'R'], uplos, t_flags, dAs, [1, 3], [1, 2], Ps
) do (side, ul, tA, dA, M, N, P)
t = tA == 'N'
R = side == 'L' ? M : N
As = blas_matrices(rng, P, R, R)
Bs = blas_matrices(rng, P, M, N)
return map(As, Bs) do A, B
(false, :none, nothing, BLAS.trmm!, side, ul, tA, dA, randn(P), A, B)
end
end...,

# trsm!
vec(
reduce(
vcat,
map(
product(['L', 'R'], ['U', 'L'], t_flags, ['N', 'U'], [1, 3], [1, 2])
) do (side, ul, tA, dA, M, N)
t = tA == 'N'
R = side == 'L' ? M : N
As = [randn(R, R) + 5I, view(randn(15, 15), 3:(R + 2), 4:(R + 3)) + 5I]
Bs = [randn(M, N), view(randn(15, 15), 2:(M + 1), 5:(N + 4))]
flags = (false, :none, nothing)
return map(product(As, Bs)) do (A, B)
(flags..., BLAS.trsm!, side, ul, tA, dA, randn(), A, B)
end
end,
),
),
map_prod(
['L', 'R'], uplos, t_flags, dAs, [1, 3], [1, 2], Ps
) do (side, ul, tA, dA, M, N, P)
t = tA == 'N'
R = side == 'L' ? M : N
As = map(blas_matrices(rng, P, R, R)) do A
A[diagind(A)] .+= 1
return A
end
Bs = blas_matrices(rng, P, M, N)
return map(As, Bs) do A, B
(false, :none, nothing, BLAS.trsm!, side, ul, tA, dA, randn(P), A, B)
end
end...,
)
memory = Any[]
return test_cases, memory
Expand Down
28 changes: 28 additions & 0 deletions src/rrules/builtins.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -744,8 +744,11 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})

# Core.Intrinsics:
(false, :stability, nothing, IntrinsicsWrappers.abs_float, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.abs_float, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.add_float, 4.0, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.add_float, 4.0f0, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.add_float_fast, 4.0, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.add_float_fast, 4.0f0, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.add_int, 1, 2),
(false, :stability, nothing, IntrinsicsWrappers.and_int, 2, 3),
(
Expand Down Expand Up @@ -793,28 +796,40 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
(false, :stability, nothing, IntrinsicsWrappers.checked_usub_int, 5, 4),
(false, :stability, nothing, IntrinsicsWrappers.copysign_float, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.copysign_float, 5.0, -3.0),
(false, :stability, nothing, IntrinsicsWrappers.copysign_float, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.copysign_float, 5.0f0, -3.0f0),
(false, :stability, nothing, IntrinsicsWrappers.ctlz_int, 5),
(false, :stability, nothing, IntrinsicsWrappers.ctpop_int, 5),
(false, :stability, nothing, IntrinsicsWrappers.cttz_int, 5),
(false, :stability, nothing, IntrinsicsWrappers.div_float, 5.0, 3.0),
(false, :stability, nothing, IntrinsicsWrappers.div_float_fast, 5.0, 3.0),
(false, :stability, nothing, IntrinsicsWrappers.div_float, 5.0f0, 3.0f0),
(false, :stability, nothing, IntrinsicsWrappers.div_float_fast, 5.0f0, 3.0f0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float, 4.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float, 4.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float_fast, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float_fast, 4.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float_fast, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.eq_float_fast, 4.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.eq_int, 5, 4),
(false, :stability, nothing, IntrinsicsWrappers.eq_int, 4, 4),
(false, :stability, nothing, IntrinsicsWrappers.flipsign_int, 4, -3),
(false, :stability, nothing, IntrinsicsWrappers.floor_llvm, 4.1),
(false, :stability, nothing, IntrinsicsWrappers.fma_float, 5.0, 4.0, 3.0),
(false, :stability, nothing, IntrinsicsWrappers.fma_float, 5.0f0, 4.0f0, 3.0f0),
(true, :stability_and_allocs, nothing, IntrinsicsWrappers.fpext, Float64, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.fpiseq, 4.1, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.fpiseq, 4.0f1, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.fptosi, UInt32, 4.1),
(false, :stability, nothing, IntrinsicsWrappers.fptoui, Int32, 4.1),
(true, :stability, nothing, IntrinsicsWrappers.fptrunc, Float32, 5.0),
(true, :stability, nothing, IntrinsicsWrappers.have_fma, Float64),
(false, :stability, nothing, IntrinsicsWrappers.le_float, 4.1, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.le_float, 4.0f1, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.le_float_fast, 4.1, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.le_float_fast, 4.0f1, 4.0f0),
# llvm_call -- NEEDS IMPLEMENTING AND TESTING
(
false,
Expand All @@ -825,17 +840,26 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
0x0000000000000018,
),
(false, :stability, nothing, IntrinsicsWrappers.lt_float, 4.1, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.lt_float, 4.0f1, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.lt_float_fast, 4.1, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.lt_float_fast, 4.0f1, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.mul_float, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.mul_float, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.mul_float_fast, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.mul_float_fast, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.mul_int, 5, 4),
(false, :stability, nothing, IntrinsicsWrappers.muladd_float, 5.0, 4.0, 3.0),
(false, :stability, nothing, IntrinsicsWrappers.muladd_float, 5.0f0, 4.0f0, 3.0f0),
(false, :stability, nothing, IntrinsicsWrappers.ne_float, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.ne_float, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.ne_float_fast, 5.0, 4.0),
(false, :stability, nothing, IntrinsicsWrappers.ne_float_fast, 5.0f0, 4.0f0),
(false, :stability, nothing, IntrinsicsWrappers.ne_int, 5, 4),
(false, :stability, nothing, IntrinsicsWrappers.ne_int, 5, 5),
(false, :stability, nothing, IntrinsicsWrappers.neg_float, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.neg_float, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.neg_float_fast, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.neg_float_fast, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.neg_int, 5),
(false, :stability, nothing, IntrinsicsWrappers.not_int, 5),
(false, :stability, nothing, IntrinsicsWrappers.or_int, 5, 5),
Expand Down Expand Up @@ -869,10 +893,14 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
(false, :stability, nothing, IntrinsicsWrappers.sle_int, 5, 4),
(false, :stability, nothing, IntrinsicsWrappers.slt_int, 4, 5),
(false, :stability, nothing, IntrinsicsWrappers.sqrt_llvm, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.sqrt_llvm, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.sqrt_llvm_fast, 5.0),
(false, :stability, nothing, IntrinsicsWrappers.sqrt_llvm_fast, 5.0f0),
(false, :stability, nothing, IntrinsicsWrappers.srem_int, 4, 1),
(false, :stability, nothing, IntrinsicsWrappers.sub_float, 4.0, 1.0),
(false, :stability, nothing, IntrinsicsWrappers.sub_float, 4.0f0, 1.0f0),
(false, :stability, nothing, IntrinsicsWrappers.sub_float_fast, 4.0, 1.0),
(false, :stability, nothing, IntrinsicsWrappers.sub_float_fast, 4.0f0, 1.0f0),
(false, :stability, nothing, IntrinsicsWrappers.sub_int, 4, 1),
(false, :stability, nothing, IntrinsicsWrappers.trunc_int, UInt8, 78),
(false, :stability, nothing, IntrinsicsWrappers.trunc_llvm, 5.1),
Expand Down
Loading

2 comments on commit 102f900

@willtebbutt
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/121191

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.4.60 -m "<description of version>" 102f900550506e457c789e79600b422f45dcabf4
git push origin v0.4.60

Please sign in to comment.