26 add optimization

.github/workflows/benchmark_pr.yml

@@ -17,7 +17,42 @@ jobs:
             - uses: julia-actions/setup-julia@v1
               with:
                 version: "1.8"
-            - uses: julia-actions/cache@v1
+            - uses: actions/cache@v2
+              name: Cache Julia packages
+              with:
+                path: ~/.julia
+                key: ${{ runner.os }}-julia-${{ hashFiles('**/Project.toml', '**/Manifest.toml') }}
+                restore-keys: |
+                  ${{ runner.os }}-julia-
+            - uses: actions/cache@v2
+              name: Cache Python packages
+              with:
+                path: ~/gt4py-venv
+                key: ${{ runner.os }}-python-${{ hashFiles('**/requirements-dev.txt') }}
+                restore-keys: |
+                  ${{ runner.os }}-python-
+            - name: Set up Python environment
+              run: |
+                sudo apt-get update
+                sudo apt-get install python3-pip python3-venv
+                python3 -m venv ~/gt4py-venv
+                source ~/gt4py-venv/bin/activate
+                python3 -m pip install --upgrade pip
+            - name: Install GT4Py from specific branch
+              run: |
+                source ~/gt4py-venv/bin/activate
+                git clone --branch fix_python_interp_path_in_cmake https://github.com/tehrengruber/gt4py.git ~/gt4py
+                cd ~/gt4py
+                pip install -r requirements-dev.txt
+                pip install -e .
+            - name: Install and Configure PyCall
+              run: |
+                source ~/gt4py-venv/bin/activate
+                julia --project=. -e 'using Pkg; Pkg.add("PyCall"); ENV["PYTHON"]="python"; Pkg.build("PyCall");'
+            - name: Check PyCall Configuration
+              run: |
+                source ~/gt4py-venv/bin/activate
+                julia --project=. -e 'using PyCall; @show PyCall.python'
             - name: Extract Package Name from Project.toml
               id: extract-package-name
               run: |
@@ -27,20 +62,20 @@ jobs:
               env:
                 JULIA_NUM_THREADS: 2
               run: |
-                # Lightweight build step, as sometimes the runner runs out of memory:
-                julia -e 'ENV["JULIA_PKG_PRECOMPILE_AUTO"]=0; import Pkg; Pkg.add(;url="https://github.com/MilesCranmer/AirspeedVelocity.jl.git")'
-                julia -e 'ENV["JULIA_PKG_PRECOMPILE_AUTO"]=0; import Pkg; Pkg.build("AirspeedVelocity")'
+                julia --project=. -e 'ENV["JULIA_PKG_PRECOMPILE_AUTO"]=0; using Pkg; Pkg.add(url="https://github.com/MilesCranmer/AirspeedVelocity.jl.git"); Pkg.build("AirspeedVelocity")'
             - name: Add ~/.julia/bin to PATH
               run: |
                 echo "$HOME/.julia/bin" >> $GITHUB_PATH
             - name: Run benchmarks
               run: |
+                source ~/gt4py-venv/bin/activate
                 echo $PATH
                 ls -l ~/.julia/bin
                 mkdir results
                 benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{github.event.repository.default_branch}}" --output-dir=results/ --tune
             - name: Create plots from benchmarks
               run: |
+                source ~/gt4py-venv/bin/activate
                 mkdir -p plots
                 benchpkgplot ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --npart=10 --format=png --input-dir=results/ --output-dir=plots/
             - name: Upload plot as artifact
@@ -50,6 +85,7 @@ jobs:
                 path: plots
             - name: Create markdown table from benchmarks
               run: |
+                source ~/gt4py-venv/bin/activate
                 benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio > table.md
                 echo '### Benchmark Results' > body.md
                 echo '' >> body.md
@@ -75,4 +111,4 @@ jobs:
                 # comment-id: ${{ steps.fcbenchmark.outputs.comment-id }}
                 issue-number: ${{ github.event.pull_request.number }}
                 body-path: body.md
-                edit-mode: replace
+                edit-mode: replace

.gitignore

@@ -18,3 +18,17 @@ docs/build/
 .DS_Store
 
 Manifest.toml
+
+# Python Env
+.venv
+env_setup.sh
+.python-version
+
+# Misc
+**/.DS_Store
+.vscode
+
+# Ignore benchmark (benchpkg) results
+results_GridTools@*
+plot_*.png
+plot_*.pdf

Project.toml

@@ -12,6 +12,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"

advection/README.md

@@ -1,6 +1,6 @@
-### README for Running `advection_miniapp.jl`
+### README for Running `advection_setup.jl` using `run_simulation_loop.jl`
 
-This README provides instructions on how to run the `advection_miniapp.jl` script for simulating advection using the Atlas library. The script allows for terminal visualization, which can be enabled as described below.
+This README provides instructions on how to run the `run_simulation_loop.jl` script for simulating advection using the Atlas library. The script allows for terminal visualization, which can be enabled as described below.
 
 #### Prerequisites
 
@@ -15,23 +15,23 @@ This README provides instructions on how to run the `advection_miniapp.jl` scrip
      ```
 
 2. **Enabling Visualization** (optional):
-   - The script has a `VISUALIZATION_FLAG` that can be set to enable or disable visualization on the terminal. Ensure that this flag is set to `true` in the `advection_miniapp.jl` script if you wish to enable visualization.
-   - Note: Other parameters such as the number of iterations can be changed in the `# Simulation Parameters` section of the script.
+   - The script has a `VISUALIZATION_FLAG` that can be set to enable or disable visualization on the terminal. Ensure that this flag is set to `true` in the `run_simulation_loop.jl` script if you wish to enable visualization.
+   - Note: Other parameters such as the number of iterations can be changed in the `# Simulation Parameters` section of the `advection_setup.jl` script.
 
 #### Running the Simulation
 
 1. **Running the Script**:
-   - Use the following command to run the `advection_miniapp.jl` script with Julia:
+   - Use the following command to run the `run_simulation_loop.jl` script with Julia:
      ```sh
-     julia --color=yes --project=$GRIDTOOLS_JL_PATH/GridTools.jl $GRIDTOOLS_JL_PATH/GridTools.jl/src/examples/advection/advection_miniapp.jl
+     julia --color=yes --project=$GRIDTOOLS_JL_PATH/GridTools.jl $GRIDTOOLS_JL_PATH/GridTools.jl/src/examples/advection/run_simulation_loop.jl
      ```
 
 #### Example
 
-Here is an example of how to set the `VISUALIZATION_FLAG` in the `advection_miniapp.jl` script and run the simulation:
+Here is an example of how to set the `VISUALIZATION_FLAG` in the `run_simulation_loop.jl` script and run the simulation:
 
 1. **Setting the Visualization Flag**:
-   - Open the `advection_miniapp.jl` script.
+   - Open the `run_simulation_loop.jl` script.
    - Set the `VISUALIZATION_FLAG` to `true`:
      ```julia
      const VISUALIZATION_FLAG = true
@@ -42,7 +42,7 @@ Here is an example of how to set the `VISUALIZATION_FLAG` in the `advection_mini
    - Run the script with the following command:
      ```sh
      export GRIDTOOLS_JL_PATH=...
-     julia --color=yes --project=. $GRIDTOOLS_JL_PATH/src/examples/advection/advection_miniapp.jl
+     julia --color=yes --project=. $GRIDTOOLS_JL_PATH/src/examples/advection/run_simulation_loop.jl
      ```
 
-By following these steps, you should be able to run the `advection_miniapp.jl` script and visualize the advection simulation results on your terminal.
+By following these steps, you should be able to run the `run_simulation_loop.jl` script and visualize the advection simulation results on your terminal.

advection/advection.jl

@@ -6,11 +6,10 @@
     level_indices::Field{Tuple{K_}, Int64},
     num_level::Int64
 )::Field{Tuple{Vertex_, K_}, Float64}
-
     return where(
-        level_indices .== num_level - 1,
+        level_indices .== 0,
         lower,
-        where(slice(level_indices .== 0, 1:29), upper, interior)
+        where(slice(level_indices .== 29, 2:30), upper, interior)
     )
 end
 
@@ -149,7 +148,8 @@ end
 )::Field{Tuple{Vertex_, K_}, Float64}
     zrhin =
         (1.0 ./ vol) .* neighbor_sum(
-            -min.(0.0, flux(V2E)) .* max.(0.0, dual_face_orientation) -
+            # TODO: fix the 0-min workaround due to the binary/unary operation issue
+            (broadcast(0., (Vertex, V2EDim, K)) .- min.(0.0, flux(V2E))) .* max.(0.0, dual_face_orientation) -
             max.(0.0, flux(V2E)) .* min.(0.0, dual_face_orientation),
             axis = V2EDim,
         )
@@ -227,15 +227,6 @@ end
     dual_face_orientation::Field{Tuple{Vertex_, V2EDim_}, Float64},
     dual_face_normal_weighted_x::Field{Tuple{Edge_}, Float64},
     dual_face_normal_weighted_y::Field{Tuple{Edge_}, Float64},
-    tmp_vertex_1::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_vertex_2::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_vertex_3::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_vertex_4::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_vertex_5::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_vertex_6::Field{Tuple{Vertex_, K_}, Float64},
-    tmp_edge_1::Field{Tuple{Edge_, K_}, Float64},
-    tmp_edge_2::Field{Tuple{Edge_, K_}, Float64},
-    tmp_edge_3::Field{Tuple{Edge_, K_}, Float64},
 )
 
     tmp_edge_1 = advector_normal(

advection/advection_miniapp.jl → advection/advection_setup.jl

@@ -1,55 +1,26 @@
-# Advection Miniapp
-# This script demonstrates an advection simulation using the Atlas library.
+# Advection Setup
+# This script demonstrates the setup of an advection simulation using the Atlas library.
 
 using Printf
-using Debugger
 using Statistics
-using Profile
 using GridTools
-
-const global VISUALIZATION_FLAG::Bool=false
-
-# Mesh Definitions --------------------------------------------------------------------------------------------
-# Define dimensions for the mesh
-Cell_ = Dimension{:Cell_, HORIZONTAL}
-Edge_ = Dimension{:Edge_, HORIZONTAL}
-Vertex_ = Dimension{:Vertex_, HORIZONTAL}
-K_ = Dimension{:K_, VERTICAL}
-V2VDim_ = Dimension{:V2V_, LOCAL}
-V2EDim_ = Dimension{:V2E_, LOCAL}
-E2VDim_ = Dimension{:E2V_, LOCAL}
-
-# Instantiate dimension objects
-Cell = Cell_()
-K = K_()
-Edge = Edge_()
-Vertex = Vertex_()
-V2VDim = V2VDim_()
-V2EDim = V2EDim_()
-E2VDim = E2VDim_()
-
-# Define field offsets to describe the relationships between different dimensions
-V2V = FieldOffset("V2V", source = Vertex, target = (Vertex, V2VDim))
-E2V = FieldOffset("E2V", source = Vertex, target = (Edge, E2VDim))
-V2E = FieldOffset("V2E", source = Edge, target = (Vertex, V2EDim))
-Koff = FieldOffset("Koff", source = K, target = K)
+using GridTools.ExampleMeshes.Unstructured
+using GridTools.AtlasMeshes
 
 # Include additional necessary files for mesh, state container, metric calculations, and advection operations
-include("../src/atlas/atlas_mesh.jl")
 include("state_container.jl")
 include("metric.jl")
 include("advection.jl")
-include("visualization_utils.jl")
 
 # Grid and Mesh Initialization --------------------------------------------------------------------------------
 # Create a structured grid and mesh for the simulation
-grid = atlas.StructuredGrid("O50")
+grid = atlas.StructuredGrid("O90")
 mesh = AtlasMesh(grid, num_level = 30)
 
 # Simulation Parameters ---------------------------------------------------------------------------------------
 δt = 1800.0  # time step in s
 niter = 50
-ε = 1.0e-8
+ϵ = 1.0e-8
 
 # Calculate metric properties from the mesh
 metric = m_from_mesh(mesh)
@@ -188,53 +159,3 @@ nabla_z(
     out = tmp_fields["tmp_vertex_2"],
     offset_provider = mesh.offset_provider
 )
-
-if VISUALIZATION_FLAG
-    # Precompute the mapping between the unstructured domain to the structured one for ASCII art visualization
-    grid_size = 50
-    mapping = precompute_mapping(mesh, xlim, ylim, grid_size)
-end
-
-# Main Simulation Loop ----------------------------------------------------------------------------------------
-for i = 1:niter
-    # Perform the upwind advection scheme to update the scalar field (rho)
-    upwind_scheme(
-        state.rho,
-        δt,
-        mesh.vol,
-        metric.gac,
-        state.vel[1],
-        state.vel[2],
-        state.vel[3],
-        mesh.pole_edge_mask,
-        mesh.dual_face_orientation,
-        mesh.dual_face_normal_weighted_x,
-        mesh.dual_face_normal_weighted_y,
-        out = state_next.rho,
-        offset_provider = mesh.offset_provider
-    )
-
-    # Print the current timestep
-    println("Timestep $i")
-
-    if VISUALIZATION_FLAG
-        # Print the current state as ASCII art every 5 timesteps
-        print_state_ascii(state, mesh, mapping, i, grid_size)
-    end
-
-    # TODO: make a function out of this switch
-    # Swap the current and next state
-    temp = state
-    global state = state_next
-    global state_next = temp
-
-    # Update the periodic boundary layers
-    update_periodic_layers(mesh, state.rho)
-end
-
-# Output the final statistics for the scalar field (rho) and velocity fields
-println(
-    "min max sum of final rho = $(minimum(state.rho.data)) , $(maximum(state.rho.data)) , $(sum(state.rho.data))"
-)
-println("Final Vel0 sum after $niter iterations: $(sum(state.vel[1].data))")
-println("Final Vel1 sum after $niter iterations: $(sum(state.vel[2].data))")

advection/run_simulation_loop.jl

@@ -0,0 +1,62 @@
+# Run Advection Miniapp Simulation
+# This script demonstrates an advection simulation using the Atlas library.
+
+include("visualization_utils.jl")
+include("advection_setup.jl")
+
+const global VISUALIZATION_FLAG::Bool=false
+const global VERBOSE_FLAG::Bool=true
+
+if VISUALIZATION_FLAG
+    # Precompute the mapping between the unstructured domain to the structured one for ASCII art visualization
+    grid_size = 50
+    mapping = precompute_mapping(mesh, xlim, ylim, grid_size)
+end
+
+# Main Simulation Loop ----------------------------------------------------------------------------------------
+for i = 1:niter
+    # Perform the upwind advection scheme to update the scalar field (rho)
+    upwind_scheme(
+        state.rho,
+        δt,
+        mesh.vol,
+        metric.gac,
+        state.vel[1],
+        state.vel[2],
+        state.vel[3],
+        mesh.pole_edge_mask,
+        mesh.dual_face_orientation,
+        mesh.dual_face_normal_weighted_x,
+        mesh.dual_face_normal_weighted_y,
+        out = state_next.rho,
+        offset_provider = mesh.offset_provider
+    )
+
+    # Print the current timestep
+    if VERBOSE_FLAG
+        println("Timestep $i")
+    end
+
+    if VISUALIZATION_FLAG
+        # Print the current state as ASCII art every 5 timesteps
+        print_state_ascii(state, mesh, mapping, i, grid_size)
+    end
+
+    # TODO: make a function out of this switch
+    # Swap the current and next state
+    temp = state
+    global state = state_next
+    global state_next = temp
+
+    # Update the periodic boundary layers
+    update_periodic_layers(mesh, state.rho)
+end
+
+if VERBOSE_FLAG
+    # Output the final statistics for the scalar field (rho) and velocity fields
+    println(
+        "min max sum of final rho = $(minimum(state.rho.data)) , $(maximum(state.rho.data)) , $(sum(state.rho.data))"
+    )
+    println("Final Vel0 sum after $niter iterations: $(sum(state.vel[1].data))")
+    println("Final Vel1 sum after $niter iterations: $(sum(state.vel[2].data))")
+end