Add macro trixi_include_changeprecision to make a double precision elixir run with single precision (#35)

efaulhaber · ranocha · web-flow · commit 9595823e0d7c · 2025-01-28T09:45:23.000+01:00
* Make `trixi_include` more flexible by allowing a mapping to be passed

* Bump version

* Update docstring of `trixi_include`

* Add `trixi_include_changeprecision`

* Export new macro

* Fix docs

* Update src/trixi_include.jl

Co-authored-by: Erik Faulhaber &lt;44124897+efaulhaber@users.noreply.github.com&gt;

* Add tests

* Reformat code

* Fix tests on Windows

* Cover missing function

* Use `mktemp` for temporary files

* Reformat

* Update src/trixi_include.jl

---------

Co-authored-by: Hendrik Ranocha &lt;ranocha@users.noreply.github.com&gt;
diff --git a/Project.toml b/Project.toml
@@ -1,9 +1,10 @@
 name = "TrixiBase"
 uuid = "9a0f1c46-06d5-4909-a5a3-ce25d3fa3284"
 authors = ["Michael Schlottke-Lakemper <michael@sloede.com>"]
-version = "0.1.5-DEV"
+version = "0.1.5"
 
 [deps]
+ChangePrecision = "3cb15238-376d-56a3-8042-d33272777c9a"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
@@ -13,6 +14,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 TrixiBaseMPIExt = "MPI"
 
 [compat]
+ChangePrecision = "1.1.0"
 MPI = "0.20"
 TimerOutputs = "0.5.25"
 julia = "1.8"
diff --git a/src/TrixiBase.jl b/src/TrixiBase.jl
@@ -1,11 +1,12 @@
 module TrixiBase
 
+using ChangePrecision: ChangePrecision
 using TimerOutputs: TimerOutput, TimerOutputs
 
 include("trixi_include.jl")
 include("trixi_timeit.jl")
 
-export trixi_include
+export trixi_include, trixi_include_changeprecision
 export @trixi_timeit, timer, timeit_debug_enabled,
        disable_debug_timings, enable_debug_timings
 
diff --git a/src/trixi_include.jl b/src/trixi_include.jl
@@ -3,7 +3,7 @@
 # of `TrixiBase`. However, users will want to evaluate in the global scope of `Main` or something
 # similar to manage dependencies on their own.
 """
-    trixi_include([mod::Module=Main,] elixir::AbstractString; kwargs...)
+    trixi_include([mapexpr::Function=identity,] [mod::Module=Main,] elixir::AbstractString; kwargs...)
 
 `include` the file `elixir` and evaluate its content in the global scope of module `mod`.
 You can override specific assignments in `elixir` by supplying keyword arguments.
@@ -16,6 +16,10 @@ into calls to `solve` with it's default value used in the SciML ecosystem
 for ODEs, see the "Miscellaneous" section of the
 [documentation](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/).
 
+The optional first argument `mapexpr` can be used to transform the included code before
+it is evaluated: for each parsed expression `expr` in `elixir`, the `include` function
+actually evaluates `mapexpr(expr)`. If it is omitted, `mapexpr` defaults to `identity`.
+
 # Examples
 
 ```@example
@@ -30,7 +34,7 @@ julia> redirect_stdout(devnull) do
 0.1
 ```
 """
-function trixi_include(mod::Module, elixir::AbstractString; kwargs...)
+function trixi_include(mapexpr::Function, mod::Module, elixir::AbstractString; kwargs...)
     # Check that all kwargs exist as assignments
     code = read(elixir, String)
     expr = Meta.parse("begin \n$code \nend")
@@ -45,13 +49,63 @@ function trixi_include(mod::Module, elixir::AbstractString; kwargs...)
     if !mpi_isparallel(Val{:MPIExt}())
         @info "You just called `trixi_include`. Julia may now compile the code, please be patient."
     end
-    Base.include(ex -> replace_assignments(insert_maxiters(ex); kwargs...), mod, elixir)
+    Base.include(ex -> mapexpr(replace_assignments(insert_maxiters(ex); kwargs...)),
+                 mod, elixir)
+end
+
+function trixi_include(mod::Module, elixir::AbstractString; kwargs...)
+    trixi_include(identity, mod, elixir; kwargs...)
 end
 
 function trixi_include(elixir::AbstractString; kwargs...)
     trixi_include(Main, elixir; kwargs...)
 end
 
+"""
+    trixi_include_changeprecision(T, [mod::Module=Main,] elixir::AbstractString; kwargs...)
+
+`include` the elixir `elixir` and evaluate its content in the global scope of module `mod`.
+You can override specific assignments in `elixir` by supplying keyword arguments,
+similar to [`trixi_include`](@ref).
+
+The only difference to [`trixi_include`](@ref) is that the precision of floating-point
+numbers in the included elixir is changed to `T`.
+More precisely, the package [ChangePrecision.jl](https://github.com/JuliaMath/ChangePrecision.jl)
+is used to convert all `Float64` literals, operations like `/` that produce `Float64` results,
+and functions like `ones` that return `Float64` arrays by default, to the desired type `T`.
+See the documentation of ChangePrecision.jl for more details.
+
+The purpose of this function is to conveniently run a full simulation with `Float32`,
+which is orders of magnitude faster on most GPUs than `Float64`, by just including
+the elixir with `trixi_include_changeprecision(Float32, elixir)`.
+Many constructors in the Trixi.jl framework are written in a way that changing all floating-point
+arguments to `Float32` will change the element type to `Float32` as well.
+In TrixiParticles.jl, including an elixir with this macro should be sufficient
+to run the full simulation with single precision.
+"""
+function trixi_include_changeprecision(T, mod::Module, filename::AbstractString; kwargs...)
+    trixi_include(expr -> ChangePrecision.changeprecision(T, replace_trixi_include(T, expr)),
+                  mod, filename; kwargs...)
+end
+
+function trixi_include_changeprecision(T, filename::AbstractString; kwargs...)
+    trixi_include_changeprecision(T, Main, filename; kwargs...)
+end
+
+function replace_trixi_include(T, expr)
+    expr = TrixiBase.walkexpr(expr) do x
+        if x isa Expr
+            if x.head === :call && x.args[1] === :trixi_include
+                x.args[1] = :trixi_include_changeprecision
+                insert!(x.args, 2, :($T))
+            end
+        end
+        return x
+    end
+
+    return expr
+end
+
 # Insert the keyword argument `maxiters` into calls to `solve` and `Trixi.solve`
 # with default value `10^5` if it is not already present.
 function insert_maxiters(expr)
diff --git a/test/trixi_include.jl b/test/trixi_include.jl
@@ -4,31 +4,27 @@
             x = 4
             """
 
-        filename = tempname()
-        try
-            open(filename, "w") do file
-                write(file, example)
-            end
+        mktemp() do path, io
+            write(io, example)
+            close(io)
 
             # Use `@trixi_testset`, which wraps code in a temporary module, and call
             # `trixi_include` with `@__MODULE__` in order to isolate this test.
-            @test_nowarn_mod trixi_include(@__MODULE__, filename)
+            @test_nowarn_mod trixi_include(@__MODULE__, path)
             @test @isdefined x
             @test x == 4
 
-            @test_nowarn_mod trixi_include(@__MODULE__, filename, x = 7)
+            @test_nowarn_mod trixi_include(@__MODULE__, path, x = 7)
 
             @test x == 7
 
             # Verify default version (that includes in `Main`)
-            @test_nowarn_mod trixi_include(filename, x = 11)
+            @test_nowarn_mod trixi_include(path, x = 11)
             @test Main.x == 11
 
             @test_throws "assignment `y` not found in expression" trixi_include(@__MODULE__,
-                                                                                filename,
+                                                                                path,
                                                                                 y = 3)
-        finally
-            rm(filename, force = true)
         end
     end
 
@@ -40,22 +36,18 @@
             x = solve()
             """
 
-        filename = tempname()
-        try
-            open(filename, "w") do file
-                write(file, example)
-            end
+        mktemp() do path, io
+            write(io, example)
+            close(io)
 
             # Use `@trixi_testset`, which wraps code in a temporary module, and call
             # `trixi_include` with `@__MODULE__` in order to isolate this test.
             @test_throws "no method matching solve(; maxiters" trixi_include(@__MODULE__,
-                                                                             filename)
+                                                                             path)
 
             @test_throws "no method matching solve(; maxiters" trixi_include(@__MODULE__,
-                                                                             filename,
+                                                                             path,
                                                                              maxiters = 3)
-        finally
-            rm(filename, force = true)
         end
     end
 
@@ -81,49 +73,112 @@
             y = solve(; maxiters=0)
             """
 
-        filename1 = tempname()
-        filename2 = tempname()
-        filename3 = tempname()
-        filename4 = tempname()
-        try
-            open(filename1, "w") do file
-                write(file, example1)
-            end
-            open(filename2, "w") do file
-                write(file, example2)
-            end
-            open(filename3, "w") do file
-                write(file, example3)
-            end
-            open(filename4, "w") do file
-                write(file, example4)
+        mktemp() do path1, io1
+            write(io1, example1)
+            close(io1)
+
+            mktemp() do path2, io2
+                write(io2, example2)
+                close(io2)
+
+                mktemp() do path3, io3
+                    write(io3, example3)
+                    close(io3)
+
+                    mktemp() do path4, io4
+                        write(io4, example4)
+                        close(io4)
+
+                        # Use `@trixi_testset`, which wraps code in a temporary module,
+                        # and call `Base.include` and `trixi_include` with `@__MODULE__`
+                        # in order to isolate this test.
+                        Base.include(@__MODULE__, path1)
+                        @test_nowarn_mod trixi_include(@__MODULE__, path2)
+                        @test @isdefined x
+                        # This is the default `maxiters` inserted by `trixi_include`
+                        @test x == 10^5
+
+                        @test_nowarn_mod trixi_include(@__MODULE__, path2, maxiters = 7)
+                        # Test that `maxiters` got overwritten
+                        @test x == 7
+
+                        # Verify that existing `maxiters` is added exactly once in the
+                        # following cases:
+                        # case 1) `maxiters` is *before* semicolon in included file
+                        @test_nowarn_mod trixi_include(@__MODULE__, path3, maxiters = 11)
+                        @test y == 11
+                        # case 2) `maxiters` is *after* semicolon in included file
+                        @test_nowarn_mod trixi_include(@__MODULE__, path3, maxiters = 14)
+                        @test y == 14
+                    end
+                end
             end
+        end
+    end
+end
+
+@trixi_testset "`trixi_include_changeprecision`" begin
+    @trixi_testset "Basic" begin
+        example = """
+            x = 4.0
+            y = zeros(3)
+            """
+
+        mktemp() do path, io
+            write(io, example)
+            close(io)
 
             # Use `@trixi_testset`, which wraps code in a temporary module, and call
-            # `Base.include` and `trixi_include` with `@__MODULE__` in order to isolate this test.
-            Base.include(@__MODULE__, filename1)
-            @test_nowarn_mod trixi_include(@__MODULE__, filename2)
+            # `trixi_include_changeprecision` with `@__MODULE__` in order to isolate this test.
+            @test_nowarn_mod trixi_include_changeprecision(Float32, @__MODULE__, path)
             @test @isdefined x
-            # This is the default `maxiters` inserted by `trixi_include`
-            @test x == 10^5
+            @test x == 4
+            @test typeof(x) == Float32
+            @test @isdefined y
+            @test eltype(y) == Float32
+
+            # Manually overwritten assignments are also changed
+            @test_nowarn_mod trixi_include_changeprecision(Float32, @__MODULE__, path,
+                                                           x = 7.0)
 
-            @test_nowarn_mod trixi_include(@__MODULE__, filename2,
-                                           maxiters = 7)
-            # Test that `maxiters` got overwritten
             @test x == 7
+            @test typeof(x) == Float32
+
+            # Verify default version (that includes in `Main`)
+            @test_nowarn_mod trixi_include_changeprecision(Float32, path, x = 11.0)
+            @test Main.x == 11
+            @test typeof(Main.x) == Float32
+        end
+    end
+
+    @trixi_testset "Recursive" begin
+        example1 = """
+            x = 4.0
+            y = zeros(3)
+            """
 
-            # Verify that adding `maxiters` to `maxiters` results in exactly one of them
-            # case 1) `maxiters` is *before* semicolon in included file
-            @test_nowarn_mod trixi_include(@__MODULE__, filename3, maxiters = 11)
-            @test y == 11
-            # case 2) `maxiters` is *after* semicolon in included file
-            @test_nowarn_mod trixi_include(@__MODULE__, filename3, maxiters = 14)
-            @test y == 14
-        finally
-            rm(filename1, force = true)
-            rm(filename2, force = true)
-            rm(filename3, force = true)
-            rm(filename4, force = true)
+        mktemp() do path1, io1
+            write(io1, example1)
+            close(io1)
+
+            # Use raw string to allow backslashes in Windows paths
+            example2 = """
+                trixi_include(@__MODULE__, raw"$path1", x = 7.0)
+                """
+
+            mktemp() do path2, io2
+                write(io2, example2)
+                close(io2)
+
+                # Use `@trixi_testset`, which wraps code in a temporary module, and call
+                # `trixi_include_changeprecision` with `@__MODULE__` in order to isolate this test.
+                @test_nowarn_mod trixi_include_changeprecision(Float32, @__MODULE__, path2)
+                @test @isdefined x
+                @test x == 7
+                @test typeof(x) == Float32
+                @test @isdefined y
+                @test eltype(y) == Float32
+            end
         end
     end
 end
