From a54ad5bb21454d3170fc8caf645877d46cbc585a Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Sun, 31 Jul 2022 18:46:34 -0700 Subject: [PATCH] Generalize axis tensor tests --- .../axistensor_conversion_benchmarks.jl | 201 ++++++++++-------- test/Geometry/ref_funcs.jl | 22 ++ test/Geometry/transform_project.jl | 116 ++++++++++ 3 files changed, 251 insertions(+), 88 deletions(-) create mode 100644 test/Geometry/ref_funcs.jl diff --git a/test/Geometry/axistensor_conversion_benchmarks.jl b/test/Geometry/axistensor_conversion_benchmarks.jl index 9c2195b963..5e727d5843 100644 --- a/test/Geometry/axistensor_conversion_benchmarks.jl +++ b/test/Geometry/axistensor_conversion_benchmarks.jl @@ -1,114 +1,139 @@ -using Test -using StaticArrays -import BenchmarkTools -import StatsBase -import OrderedCollections -import Random +using Test, StaticArrays +import Random, BenchmarkTools, StatsBase, OrderedCollections, LinearAlgebra #! format: off -using ClimaCore.Geometry: - AbstractAxis, - CovariantAxis, - AxisVector, - ContravariantAxis, - LocalAxis, - CartesianAxis, - AxisTensor, - Covariant1Vector, - Covariant13Vector, - UVVector, - UWVector, - UVector, - WVector, - Covariant12Vector, - UVWVector, - Covariant123Vector, - Covariant3Vector, - Contravariant12Vector, - Contravariant3Vector, - Contravariant123Vector, - Contravariant13Vector, - Contravariant2Vector, - Axis2Tensor +using ClimaCore.Geometry:Geometry, AbstractAxis, CovariantAxis, + AxisVector, ContravariantAxis, LocalAxis, CartesianAxis, AxisTensor, + Covariant1Vector, Covariant13Vector, UVVector, UWVector, UVector, + WVector, Covariant12Vector, UVWVector, Covariant123Vector, Covariant3Vector, + Contravariant12Vector, Contravariant3Vector, Contravariant123Vector, + Contravariant13Vector, Contravariant2Vector, Axis2Tensor, Contravariant3Axis, + LocalGeometry, CovariantTensor, CartesianTensor, LocalTensor, ContravariantTensor include("transform_project.jl") # compact, generic but unoptimized reference include("used_transform_args.jl") +include("ref_funcs.jl") include("used_project_args.jl") -function benchmark_axistensor_conversions(args, ::Type{FT}, func::F) where {FT, F} - results = OrderedCollections.OrderedDict() - for (aTo, x) in args - func(aTo, x) # compile first - - # Setting up a benchmark is _really_ slow - # for many many microbenchmarks. So, let's - # just hardcode something simple that runs - # fast.. - - #### Using BenchmarkTools - - # b = BenchmarkTools.@benchmarkable $func($aTo, $x) - # trial = BenchmarkTools.run(b, samples = 3) - # time = StatsBase.mean(trial.times) - # time = BenchmarkTools.@btime begin; result = func($aTo, $x); end - - #### Hard-code average of 25 calls - time = @elapsed begin - func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x) - func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x) - func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x) - func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x) - func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x); func(aTo, x) - end - ns = 1e9 - time = time/25*ns # average and convert to ns +# time_func(func::F, x, y) where {F} = time_func_accurate(func, x, y) +time_func(func::F, x, y) where {F} = time_func_fast(func, x, y) + +function time_func_accurate(func::F, x, y) where {F} + b = BenchmarkTools.@benchmarkable $func($x, $y) + trial = BenchmarkTools.run(b; samples = 3) + # show(stdout, MIME("text/plain"), trial) + time = StatsBase.mean(trial.times) + return time +end + +function time_func_fast(func::F, x, y) where {F} + time = 0 + nsamples = 100 + for i in 1:nsamples + time += @elapsed func(x, y) + end + ns = 1e9 + time = time/nsamples*ns # average and convert to ns + return time +end - key = typeof.((aTo, x)) - result = func(aTo, x) +function benchmark_conversion(arg_set, ::Type{FT}, func::F) where {FT, F} + results = OrderedCollections.OrderedDict() + ref_func = reference_func(func) + for args in arg_set + key = typeof.(args) + # Reference + result = ref_func(args...) # compile first + time = time_func(ref_func, args...) t_pretty = BenchmarkTools.prettytime(time) - @info "Benchmarking $t_pretty $func with $key" - results[key] = (time, t_pretty, result) + ref = (;time, t_pretty, result) + # Optimized + result = func(args...) # compile first + time = time_func(func, args...) + t_pretty = BenchmarkTools.prettytime(time) + opt = (;time, t_pretty, result) + + @info "Benchmark: opt: $(opt.t_pretty) ref: $(ref.t_pretty). Key: $key" + results[key] = (;opt, ref) end return results end -function get_conversion_args(ucs) - map(ucs) do (axt, axtt) +function all_axes() + all_Is() = [(1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)] + collect(Iterators.flatten(map(all_Is()) do I + ( + CovariantAxis{I}(), + ContravariantAxis{I}(), + LocalAxis{I}(), + CartesianAxis{I}() + ) + end)) +end + +all_observed_axistensors(::Type{FT}) where {FT} = + vcat(map(x-> rand(last(x)), used_project_arg_types(FT)), + map(x-> rand(last(x)), used_transform_arg_types(FT))) + +func_args(FT, ::typeof(Geometry.project)) = + map(used_project_arg_types(FT)) do (axt, axtt) + (axt(), rand(axtt)) + end +func_args(FT, ::typeof(Geometry.transform)) = + map(used_transform_arg_types(FT)) do (axt, axtt) (axt(), rand(axtt)) end + +function all_possible_func_args(FT, ::typeof(Geometry.contravariant3)) + # TODO: this is not accurate yet, since we don't yet + # vary over all possible LocalGeometry's. + M = @SMatrix [ + FT(4) FT(1) + FT(0.5) FT(2) + ] + J = LinearAlgebra.det(M) + ∂x∂ξ = rand(Geometry.AxisTensor{FT, 2, Tuple{Geometry.LocalAxis{(3,)}, Geometry.CovariantAxis{(3,)}}, SMatrix{1, 1, FT, 1}}) + lg = Geometry.LocalGeometry(Geometry.XYPoint(FT(0), FT(0)), J, J, ∂x∂ξ) + # Geometry.LocalGeometry{(3,), Geometry.ZPoint{FT}, FT, SMatrix{1, 1, FT, 1}} + Iterators.flatten( + map(used_project_arg_types(FT)) do (axt, axtt) + map(all_axes()) do ax + (rand(axtt), lg) + end + end + ) end -function test_optimized_transform(::Type{FT}) where {FT} - @info "Testing optimized transform..." - utat = used_transform_arg_types(FT) - args = get_conversion_args(utat) - rt_results = benchmark_axistensor_conversions(args, FT, ref_transform) - ot_results = benchmark_axistensor_conversions(args, FT, ref_transform) # (opt_transform) - for key in keys(rt_results) - @test last(ot_results[key]) == last(rt_results[key]) # test correctness - @test_broken first(ot_results[key])*30 < first(rt_results[key]) # test performance +function func_args(FT, f::typeof(Geometry.contravariant3)) + # TODO: fix this.. + apfa = all_possible_func_args(FT, f) + args_dict = Dict() + for _args in apfa + hasmethod(f, typeof(_args)) || continue + args_dict[typeof.(_args)] = _args end + return values(args_dict) end -function test_optimized_project(::Type{FT}) where {FT} - @info "Testing optimized project..." - upat = used_project_arg_types(FT) - args = get_conversion_args(upat) - rp_results = benchmark_axistensor_conversions(args, FT, ref_project) - op_results = benchmark_axistensor_conversions(args, FT, ref_project) # (opt_project) - for key in keys(rp_results) - @test last(op_results[key]) == last(rp_results[key]) # test correctness - @test_broken first(op_results[key])*30 < first(rp_results[key]) # test performance +reference_func(::typeof(Geometry.contravariant3)) = ref_contravariant3 +reference_func(::typeof(Geometry.project)) = ref_project +reference_func(::typeof(Geometry.transform)) = ref_transform + +function test_optimized_function(::Type{FT}, func) where {FT} + @info "Testing optimized $func..." + args = func_args(FT, func) + bm = benchmark_conversion(args, FT, func) + for key in keys(bm) + @test bm[key].opt.result == bm[key].ref.result # test correctness + @test_broken bm[key].opt.time*10 < bm[key].ref.time # test performance end end -# TODO: figure out how to make error checking in `transform` - -# @testset "Test optimized transform" begin -# test_optimized_transform(Float64) -# end +# TODO: figure out how to make error checking in `transform` optional -@testset "Test optimized project" begin - test_optimized_project(Float64) +@testset "Test optimized functions" begin + test_optimized_function(Float64, Geometry.project) + # test_optimized_function(Float64, Geometry.contravariant3) + # test_optimized_function(Float64, Geometry.transform) end #! format: on diff --git a/test/Geometry/ref_funcs.jl b/test/Geometry/ref_funcs.jl new file mode 100644 index 0000000000..65397c4265 --- /dev/null +++ b/test/Geometry/ref_funcs.jl @@ -0,0 +1,22 @@ + +@inline ref_contravariant1(u::AxisVector, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant1Axis(), u, local_geometry)[1] +@inline ref_contravariant2(u::AxisVector, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant2Axis(), u, local_geometry)[1] +@inline ref_contravariant3(u::AxisVector, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant3Axis(), u, local_geometry)[1] + +@inline ref_contravariant1(u::Axis2Tensor, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant1Axis(), u, local_geometry)[1, :] +@inline ref_contravariant2(u::Axis2Tensor, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant2Axis(), u, local_geometry)[1, :] +@inline ref_contravariant3(u::Axis2Tensor, local_geometry::LocalGeometry) = + @inbounds ref_project(Contravariant3Axis(), u, local_geometry)[1, :] + +@inline ref_covariant1(u::AxisVector, local_geometry::LocalGeometry) = + CovariantVector(u, local_geometry).u₁ +@inline ref_covariant2(u::AxisVector, local_geometry::LocalGeometry) = + CovariantVector(u, local_geometry).u₂ +@inline ref_covariant3(u::AxisVector, local_geometry::LocalGeometry) = + CovariantVector(u, local_geometry).u₃ + diff --git a/test/Geometry/transform_project.jl b/test/Geometry/transform_project.jl index e8270753e4..ff1bebb091 100644 --- a/test/Geometry/transform_project.jl +++ b/test/Geometry/transform_project.jl @@ -161,3 +161,119 @@ end SMatrix{$(length(Ito)), $M}($(vals...)), )) end + + +for op in (:ref_transform, :ref_project) + @eval begin + # Covariant <-> Cartesian + @inline $op( + ax::CartesianAxis, + v::CovariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂ξ∂x' * + $op(Geometry.dual(axes(local_geometry.∂ξ∂x, 1)), v), + ) + @inline $op( + ax::CovariantAxis, + v::CartesianTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂x∂ξ' * + $op(Geometry.dual(axes(local_geometry.∂x∂ξ, 1)), v), + ) + @inline $op( + ax::LocalAxis, + v::CovariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂ξ∂x' * + $op(Geometry.dual(axes(local_geometry.∂ξ∂x, 1)), v), + ) + @inline $op( + ax::CovariantAxis, + v::LocalTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂x∂ξ' * + $op(Geometry.dual(axes(local_geometry.∂x∂ξ, 1)), v), + ) + + # Contravariant <-> Cartesian + @inline $op( + ax::ContravariantAxis, + v::CartesianTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂ξ∂x * + $op(Geometry.dual(axes(local_geometry.∂ξ∂x, 2)), v), + ) + @inline $op( + ax::CartesianAxis, + v::ContravariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂x∂ξ * + $op(Geometry.dual(axes(local_geometry.∂x∂ξ, 2)), v), + ) + @inline $op( + ax::ContravariantAxis, + v::LocalTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂ξ∂x * + $op(Geometry.dual(axes(local_geometry.∂ξ∂x, 2)), v), + ) + + @inline $op( + ax::LocalAxis, + v::ContravariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂x∂ξ * + $op(Geometry.dual(axes(local_geometry.∂x∂ξ, 2)), v), + ) + + # Covariant <-> Contravariant + @inline $op( + ax::ContravariantAxis, + v::CovariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂ξ∂x * + local_geometry.∂ξ∂x' * + $op(Geometry.dual(axes(local_geometry.∂ξ∂x, 1)), v), + ) + @inline $op( + ax::CovariantAxis, + v::ContravariantTensor, + local_geometry::LocalGeometry, + ) = $op( + ax, + local_geometry.∂x∂ξ' * + local_geometry.∂x∂ξ * + $op(Geometry.dual(axes(local_geometry.∂x∂ξ, 2)), v), + ) + + @inline $op(ato::CovariantAxis, v::CovariantTensor, ::LocalGeometry) = + $op(ato, v) + @inline $op( + ato::ContravariantAxis, + v::ContravariantTensor, + ::LocalGeometry, + ) = $op(ato, v) + @inline $op(ato::CartesianAxis, v::CartesianTensor, ::LocalGeometry) = + $op(ato, v) + @inline $op(ato::LocalAxis, v::LocalTensor, ::LocalGeometry) = + $op(ato, v) + end +end