[WIP] Enzyme Test

CMakeLists.txt

@@ -163,6 +163,10 @@ set(_ALL_TARGETS lib)
 set_target_properties(lib PROPERTIES
     POSITION_INDEPENDENT_CODE ON
     WINDOWS_EXPORT_ALL_SYMBOLS ON
+    # Only needed for Clang Plugin of Enzyme
+    UNITY_BUILD ON
+    UNITY_BUILD_MODE BATCH
+    UNITY_BUILD_BATCH_SIZE 1000
 )
 
 # executable application

README.md

@@ -1,154 +1,116 @@
-# ImpactX
+# Differentiable ImpactX
 
-[![CI Status](https://github.com/ECP-WarpX/impactx/actions/workflows/ubuntu.yml/badge.svg)](https://github.com/ECP-WarpX/impactx/actions/workflows/ubuntu.yml)
-[![Documentation Status](https://readthedocs.org/projects/impactx/badge/?version=latest)](https://impactx.readthedocs.io)
-[![License ImpactX](https://img.shields.io/badge/license-BSD--3--Clause--LBNL-blue.svg)](https://spdx.org/licenses/BSD-3-Clause-LBNL.html)
-[![Supported Platforms](https://img.shields.io/badge/platforms-linux%20|%20osx%20|%20win-blue)](https://impactx.readthedocs.io/en/latest/install/users.html)  
-[![DOI (source)](https://img.shields.io/badge/DOI%20(source)-10.5281/zenodo.6954922-blue.svg)](https://doi.org/10.5281/zenodo.6954922)
-[![DOI (paper)](https://img.shields.io/badge/DOI%20(paper)-10.18429%2FJACoW--NAPAC2022--TUYE2-blue.svg)](https://doi.org/10.18429/JACoW-NAPAC2022-TUYE2)  
-[![Language: C++17](https://img.shields.io/badge/language-C%2B%2B17-orange.svg)](https://isocpp.org/)
-[![Language: Python](https://img.shields.io/badge/language-Python-orange.svg)](https://python.org/)
 
-ImpactX: an s-based beam dynamics code including space charge effects.
-This is the next generation of the [IMPACT-Z](https://github.com/impact-lbl/IMPACT-Z) code.
+## Prepare Developer Environment (Ubuntu 24.02 w/ Clang/LLVM 16)
 
-## Documentation
-
-In order to learn how to install and run the code, please see the online documentation:
-https://impactx.readthedocs.io
-
-* ImpactX Doxygen: https://impactx.readthedocs.io/en/latest/_static/doxyhtml
-* AMReX Doxygen: https://amrex-codes.github.io/amrex/doxygen
-* WarpX Doxygen: https://warpx.readthedocs.io/en/latest/_static/doxyhtml
-
-## Contributing
+### Once
+```bash
+# Ref: https://github.com/EnzymeAD/enzyme-dev-docker/blob/main/Dockerfile
+sudo apt install llvm-dev clang-16 lld-16 zlib1g-dev libzstd-dev git automake autoconf cmake make lldb-16 ninja-build build-essential libtool llvm-16-dev libclang-16-dev libomp-16-dev libblas-dev libeigen3-dev libboost-dev python3
+python3 -m pip install --break-system-package pipx pathlib2
+python3 -m pipx install lit cmake
+
+export CC="clang-16"
+export CXX="clang++-16"
+
+# Ref: https://enzyme.mit.edu/Installation/
+git clone https://github.com/EnzymeAD/Enzyme ~/src/Enzyme
+cd ~/src/Enzyme
+cmake --fresh -G Ninja -S enzyme -B build -DLLVM_DIR=/usr/lib/llvm-16/lib/cmake/llvm -DLLVM_EXTERNAL_LIT=$(which lit)
+cmake --build build -j 6
+
+cd build
+ninja check-enzyme
+
+# linker and lld and lto enforcement in CMake is hard...
+cd ..
+mkdir mylld
+cd mylld
+ln -s /usr/lib/llvm-16/bin/lld-link lld
+ln -s $(which ld.lld-16) ld.lld  # note: not sufficient yet... somehow hard-coded in compiler detection... use docker
+                                 # manually linking /usr/bin/ld.lld-16 as /usr/bin/ld.lld works...
+export PATH=$PWD:$PATH
+```
 
-[![AMReX](https://img.shields.io/static/v1?label="runs%20on"&message="AMReX"&color="blueviolet")](https://amrex-codes.github.io/)
 
-Our workflow is described in [CONTRIBUTING.rst](CONTRIBUTING.rst).
+## Compile Options
 
-## Developer Environment
+From the [homepage](https://enzyme.mit.edu/getting_started/UsingEnzyme/):
+> Enzyme supports differentiating C/C++ code through ClangEnzyme and LLDEnzyme as compiler and linker plugins, respectively.
+> Clang gives our plugin more flexibility in adding and ordering optimization passes than LLD and therefore using ClangEnzyme could result in better performance than LLDEnzyme.
+> However, ClangEnzyme can only differentiate one compilation unit at a time and will therefore fail if the function which you try to differentiate calls functions in other compilation units (generally other .c or .cpp files).
+> In these cases we recommend the use of LLDEnzyme in combination with LTO.
 
-Please see our [developer installation section](https://impactx.readthedocs.io/en/latest/install/dependencies.html#install-dependencies) of the documentation for an easy install of our software dependencies.
 
-## Get the Source Code
+## Compile (ClangEnzyme, one TU)
 
-Before you start, you will need a copy of the ImpactX source code:
+### Always
 
 ```bash
-git clone [email protected]:ECP-WarpX/impactx.git
-cd impactx
-```
+export PATH=$HOME/.local/pipx/venvs/cmake/bin:$PATH
+export CC="clang-16"
+export CXX="clang++-16"
 
-## Compile
+# one TU: Clang Plugin
+#   AMReX globals: https://enzyme.mit.edu/getting_started/UsingEnzyme/#assume-inactivity-of-unmarked-globals
+export CXXFLAGS="-fplugin=$HOME/src/Enzyme/build/Enzyme/ClangEnzyme-16.so -mllvm -enzyme-globals-default-inactive=1"
+```
 
+With the active developer env above, inside the ImpactX source dir:
 ```bash
-# find dependencies & configure
-cmake -S . -B build
-
-# compile
-cmake --build build -j 4
+cmake --fresh \
+  -S .        \
+  -B build    \
+  -DImpactX_MPI=OFF       \
+  -DImpactX_COMPUTE=NOACC \
+  -DImpactX_OPENPMD=OFF   \
+  -DCMAKE_LINKER_TYPE=LLD \
+  -DCMAKE_LINKER=/usr/lib/llvm-16/bin/lld-link
+
+cmake --build build -j 6
 ```
 
-That's all!
-ImpactX binaries are now in `build/bin/`.
-Most people execute these binaries directly or copy them out.
 
-You can inspect and modify build options after running `cmake -S . -B` build with either
+## Compile (LLDEnzyme, multiple TUs)
 
-```bash
-ccmake build
-```
+TODO: redo this part using https://github.com/EnzymeAD/enzyme-dev-docker because ld, lld, ld.ldd with the non-system default is too tricky to get right in CMake (i.e. compiler detection).
 
-or by adding arguments with `-D<OPTION>=<VALUE>` to the first CMake call, e.g.:
+### Always
 
 ```bash
-cmake -S . -B build -DImpactX_COMPUTE=CUDA -DImpactX_MPI=OFF
+export PATH=$HOME/.local/pipx/venvs/cmake/bin:$PATH
+export CC="clang-16"
+export CXX="clang++-16"
+
+# many TU: LDD Plugin
+#   https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/Enzyme/CMakeLists.txt
+#   AMReX globals: https://enzyme.mit.edu/getting_started/UsingEnzyme/#assume-inactivity-of-unmarked-globals
+export CXXFLAGS="-fuse-ld=/usr/lib/llvm-16/bin/lld-link -flto"  # -mllvm -enzyme-globals-default-inactive=1
+export LDFLAGS="-fuse-ld=/usr/lib/llvm-16/bin/lld-link -flto -Wl,-mllvm -Wl,-load=$HOME/src/Enzyme/build/Enzyme/LLDEnzyme-16.so -Wl,--load-pass-plugin=$HOME/src/Enzyme/build/Enzyme/LLDEnzyme-16.so"
 ```
 
-### Python Compile
-
+With the active developer env above, inside the ImpactX source dir:
 ```bash
-# find dependencies & configure
-cmake -S . -B build -DImpactX_PYTHON=ON
-
-# compile & install
-cmake --build build -j 4 --target pip_install
+cmake --fresh \
+  -S .        \
+  -B build    \
+  -DImpactX_MPI=OFF       \
+  -DImpactX_COMPUTE=NOACC \
+  -DImpactX_OPENPMD=OFF   \
+  -DCMAKE_LINKER_TYPE=LLD \
+  -DCMAKE_LINKER=/usr/lib/llvm-16/bin/lld-link
+
+# note: -DCMAKE_LINKER_TYPE=LLD appends general lld not necessarily the right version
+#        bend ldd to ldd-16 and /usr/bin/ld.ldd to ld.ldd-16
+# -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON
+
+cmake --build build -j 6
 ```
 
-## Run
-
-An executable ImpactX binary with the current compile-time options encoded in its file name will be created in `build/bin/`.
-
-Additionally, a symbolic link named `impactx` can be found in that directory, which points to the last built ImpactX executable.
-
-The command-line syntax for this executable is:
-```console
-Usage: impactx <inputs-file> [some.overwritten.option=value]...
-
-Mandatory arguments (remove the <>):
-  inputs-file     the path to an input file; can be relative to the current
-                  working directory or absolute.
-                  Example: input_fodo.in
-
-Optional arguments (remove the []):
-  options         this can overwrite any line in an inputs-file
-                  Example: quad1.ds=0.5 sbend1.rc=1.5
 
-Examples:
-  In the current working directory, there is a file "input_fodo.in" and the
-  "impactx" executable.
-  The line to execute would look like this:
-    ./impactx input_fodo.in
+## Test/Run
 
-  In the current working directory, there is a file "input_fodo.in" and the
-  executable "impactx" is in a directory that is listed in the "PATH"
-  environment variable.
-  The line to execute would look like this:
-    impactx input_fodo.in
-
-  In the current working directory, there is a file "input_fodo.in" and the
-  "impactx" executable. We want to voerwrite the segment length of the beamline
-  element "quad1" that is already defined in it. We also want to change the
-  radius of curvature of the bending magnet "sbend1" to a different value than
-  in the file "input_fodo.in".
-  The line to execute would look like this:
-    ./impactx input_fodo.in quad1.ds=0.5 sbend1.rc=1.5
-```
-
-## Test
-
-In order to run our tests, you need to have a few Python packages installed:
-```console
-python3 -m pip install --upgrade pip
-python3 -m pip install --upgrade build packaging setuptools wheel pytest
-python3 -m pip install --upgrade -r tests/python/requirements.txt
-```
-
-You can run all our tests with:
-
-```console
-ctest --test-dir build --output-on-failure
+With the active developer env above:
+```bash
+./build/bin/impactx examples/fodo/input_fodo.in
 ```
-
-Further options:
-* help: `ctest --test-dir build --help`
-* list all tests: `ctest --test-dir build -N`
-* only run tests that have "FODO" in their name: `ctest --test-dir build -R FODO`
-
-## Acknowledgements
-
-This work was supported by the Laboratory Directed Research and Development Program of Lawrence Berkeley National Laboratory under U.S. Department of Energy Contract No. DE-AC02-05CH11231.
-
-ImpactX is supported by the CAMPA collaboration, a project of the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research and Office of High Energy Physics, Scientific Discovery through Advanced Computing (SciDAC) program.
-
-## Copyright Notice
-
-ImpactX Copyright (c) 2022, The Regents of the University of California, through Lawrence Berkeley National Laboratory (subject to receipt of any required approvals from the U.S. Dept. of Energy).
-All rights reserved.
-
-If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at [email protected].
-
-Please see the full license agreement in [LICENSE.txt](LICENSE.txt).  
-Please see the notices in [NOTICE.txt](NOTICE.txt).  
-The SPDX license identifier is `BSD-3-Clause-LBNL`.

cmake/dependencies/ABLASTR.cmake

@@ -60,6 +60,8 @@ macro(find_ablastr)
         set(ABLASTR_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE)
         set(AMReX_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE)
 
+        set(AMReX_TINY_PROFILE OFF CACHE BOOL "")
+
         set(WarpX_APP OFF CACHE BOOL "" FORCE)
         set(WarpX_LIB OFF CACHE BOOL "" FORCE)
         set(WarpX_QED OFF CACHE BOOL "" FORCE)

src/ImpactX.H

@@ -147,6 +147,16 @@ namespace impactx
         bool m_grids_initialized = false;
     };
 
+    extern amrex::ParticleReal __enzyme_autodiff(void*, amrex::ParticleReal);
+
+    amrex::ParticleReal
+    evaluate_lattice (amrex::ParticleReal q1_k);  //, amrex::ParticleReal q2_k);
+
+    void my_run ();
+
+    // global var until functors work better
+    static ImpactX* active_sim = nullptr;
+
 } // namespace impactx
 
 #endif // IMPACT_X_H

src/ImpactX.cpp

@@ -13,7 +13,9 @@
 #include "particles/CollectLost.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
+#include "particles/elements/All.H"
 #include "particles/diagnostics/DiagnosticOutput.H"
+#include "particles/diagnostics/ReducedBeamCharacteristics.H"
 #include "particles/spacecharge/ForceFromSelfFields.H"
 #include "particles/spacecharge/GatherAndPush.H"
 #include "particles/spacecharge/PoissonSolve.H"
@@ -322,9 +324,9 @@ namespace impactx {
                 std::string openpmd_backend = "default";
                 pp_diag.queryAdd("backend", openpmd_backend);
 
-                diagnostics::BeamMonitor output_lost("particles_lost", openpmd_backend, "g");
-                output_lost(*amr_data->m_particles_lost, 0, 0);
-                output_lost.finalize();
+                //diagnostics::BeamMonitor output_lost("particles_lost", openpmd_backend, "g");
+                //output_lost(*amr_data->m_particles_lost, 0, 0);
+                //output_lost.finalize();
             }
         }
 
@@ -336,4 +338,79 @@ namespace impactx {
             }, element_variant);
         }
     }
+
+    amrex::ParticleReal
+    evaluate_lattice (amrex::ParticleReal q1_k)  //, amrex::ParticleReal q2_k)
+    {
+        amrex::ParticleReal q2_k = 3.0;
+
+        // ns = 10  // TODO: number of slices per ds in the element
+
+        auto dr1 = Drift{2.7};
+        auto q1 = Quad{0.1, q1_k};
+        auto dr2 = Drift{1.4};
+        auto q2 = Quad{0.2, q2_k};
+        auto dr3 = Drift{1.4};
+        auto q3 = Quad{0.1, q1_k};
+        auto dr4 = Drift{2.7};
+        /*
+        // quadrupole triplet
+        // https://impactx.readthedocs.io/en/latest/usage/examples/optimize_triplet/README.html
+        active_sim->m_lattice = {
+            Drift{2.7},
+            Quad{0.1, q1_k},
+            Drift{1.4},
+            Quad{0.2, q2_k},
+            Drift{1.4},
+            Quad{0.1, q1_k},
+            Drift{2.7}
+        };
+
+
+        // int ns = 25;  // number of slices per ds in the element
+        // for (auto & e : sim.m_lattice) { e.m_nslice = ns; }
+        */
+
+        //active_sim->evolve();
+        auto & pc = *active_sim->amr_data->m_particle_container;
+
+        dr1(pc, 0, 0);
+
+        /*
+        std::unordered_map<std::string, amrex::ParticleReal> const rbc =
+                diagnostics::reduced_beam_characteristics(*active_sim->amr_data->m_particle_container);
+
+        return rbc.at("alpha_x");  // TOOD: alpha_x, alpha_y, beta_x, beta_y
+        */
+        return 0.0;
+    }
+
+    void my_run ()
+    {
+        ImpactX sim;
+
+        sim.init_grids();
+
+        // TODO: replace with beam params from https://impactx.readthedocs.io/en/latest/usage/examples/optimize_triplet/README.html
+        sim.initBeamDistributionFromInputs();
+
+        // design the accelerator lattice
+        // sim.initLatticeElementsFromInputs();
+
+        // initial quad strengths
+        amrex::ParticleReal q1_k = -3.0;
+        //amrex::ParticleReal q2_k = 3.0;
+
+        active_sim = &sim;
+
+        // non-differentiable run:
+        amrex::ParticleReal const alpha_x = evaluate_lattice(q1_k);
+        amrex::Print() << "final alpha_x = " << alpha_x << std::endl;
+
+        // differentiable run:
+        double ddx = __enzyme_autodiff((void*) evaluate_lattice, q1_k);
+        amrex::Print() << "ddx = " << ddx << std::endl;
+
+        sim.finalize();
+    }
 } // namespace impactx

src/initialization/InitElement.cpp

@@ -500,7 +500,7 @@ namespace detail
                 pp_element.queryAddWithParser("cn", cn);
             }
 
-            m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding, period_sample_intervals));
+            //m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding, period_sample_intervals));
         } else if (element_type == "line")
         {
             // Parse the lattice elements for the sub-lattice in the line