Merge changes made for Frontier (#934)

Big merge to bring all changes made for Frontier back to the main branch. **WARNING: This merge requires Kokkos 4.0 or later!** - switch to Kokkos 4.0 - bring nxmx_writer to nanoBragg - reduce memory footprint of hopper and stage2 - optimizations for nanoBragg and diffBragg stage 2 **Changes for Kokkos 4.0** - moved mathematical functions (`exp`, `sin`, ...) and `HIPSpace` from `Kokkos::experimental` to `Kokkos` namespace - Kokkos is now build with C++17 standard - Switch from UVM to manual memory management - InitArguments is deprecated, using InitializationSettings now **Changes for nxmx_writer** - Add nxmx_writer capability to nanoBragg - Uses dxtbx's general purpose nxmx_writer - Need to set values for material and thickness for NXmx - Hook up some missing pipes for ImageSetData and FormatBraggInMemory so they can work with imageset.get_spectrum - Add example test to tst_gauss_argchk.py - Work in progress writing simulations to NeXuS **Other Changes** - change `use_cuda` to `use_gpu` - added transfer function for boolean vector to kokkos view - micro optimizations change division to multiplication with reciprocal - collect individual refinement flags into one bit-mask - collect gradients in kokkos_manager struct - port rotated axis for diffuse scattering to diffbragg_kokkos - update unit tests to check that curve fit goes through origin - change panel vectors to use kokkos_vec3 instead of double[3] - in nanoBragg_kokkos, precalculate crystal_orientation in separate function - Fix hopper_utils.py to not keep full background after ROIs are created - Always use timestamps in diffBragg/mpi_logger - adds 3-fold Flatt support to diffBragg - better timers for device code - fixes to command line, notably integrate OOM error - better management of host-dev copies - use stride when spectrum_from_imageset=True --------- Co-authored-by: Felix Wittwer <[email protected]> Co-authored-by: Aaron Brewster <[email protected]> Co-authored-by: Nicholas K Sauter <[email protected]> Co-authored-by: Vidya Ganapati <[email protected]> Co-authored-by: Billy K. Poon <[email protected]> Co-authored-by: Daniel Tchoń <[email protected]> Co-authored-by: Michael Wall <[email protected]> Co-authored-by: dermen <[email protected]>
cctbx · Jan 8, 2024 · 064ef40 · 064ef40
1 parent 432483d
commit 064ef40
Show file tree

Hide file tree

Showing 83 changed files with 5,405 additions and 2,352 deletions.
diff --git a/.azure-pipelines/xfel/download-source.yml b/.azure-pipelines/xfel/download-source.yml
@@ -61,7 +61,9 @@ jobs:
   # preserve permissions and delete extra files
   - script: |
       cd $(Pipeline.Workspace)
+      mv modules/kokkos/.git modules/kokkos/.git.tmp
       rm -fr modules/*/.git/*
+      mv modules/kokkos/.git.tmp modules/kokkos/.git
       rm -fr modules/*/.svn/*
       rm -fr modules/*.tar
       tar -cf modules.tar modules

diff --git a/.azure-pipelines/xfel/unix-conda-build.yml b/.azure-pipelines/xfel/unix-conda-build.yml
@@ -95,6 +95,7 @@ steps:
       libtbx.configure lunus
       make
     fi
+    rm -fr modules/kokkos/.git/*
   displayName: Configure and Build
 
 # test

diff --git a/kokkostbx/Legacy/kokkos_matrix.h b/kokkostbx/Legacy/kokkos_matrix.h
@@ -156,7 +156,7 @@ namespace kokkostbx {
         }
 
         KOKKOS_INLINE_FUNCTION NumType length() const {
-            return ::Kokkos::Experimental::sqrt(length_sqr());
+            return ::Kokkos::sqrt(length_sqr());
         }
 
         KOKKOS_INLINE_FUNCTION NumType dot(const matrix<NumType>& v) const {

diff --git a/kokkostbx/Legacy/kokkos_matrix3.h b/kokkostbx/Legacy/kokkos_matrix3.h
@@ -147,7 +147,7 @@ namespace kokkostbx {
         }
 
         KOKKOS_INLINE_FUNCTION NumType length() const {
-            return ::Kokkos::Experimental::sqrt(length_sqr());
+            return ::Kokkos::sqrt(length_sqr());
         }
 
         KOKKOS_INLINE_FUNCTION NumType dot(const matrix3<NumType>& v) const {

diff --git a/kokkostbx/Legacy/kokkos_vector.h b/kokkostbx/Legacy/kokkos_vector.h
@@ -215,7 +215,7 @@ namespace kokkostbx {
         }
 
         KOKKOS_INLINE_FUNCTION NumType length() const {
-            return ::Kokkos::Experimental::sqrt(length_sqr());
+            return ::Kokkos::sqrt(length_sqr());
         }
 
         KOKKOS_INLINE_FUNCTION NumType dot(const vector3<NumType>& v) const {
@@ -248,8 +248,8 @@ namespace kokkostbx {
 
         // rotate a point about a unit vector3 axis
         KOKKOS_INLINE_FUNCTION vector3<NumType> rotate_around_axis(const vector3<NumType>& axis, NumType angle) const {
-            NumType sinphi = ::Kokkos::Experimental::sin(angle);
-            NumType cosphi = ::Kokkos::Experimental::cos(angle);
+            NumType sinphi = ::Kokkos::sin(angle);
+            NumType cosphi = ::Kokkos::cos(angle);
             NumType dot_factor = axis.dot(*this) * (1.0-cosphi);
 
             vector3<NumType> vector_rot = axis.cross(*this) * sinphi;

diff --git a/kokkostbx/Legacy/kokkos_vector3.h b/kokkostbx/Legacy/kokkos_vector3.h
@@ -144,7 +144,7 @@ namespace kokkostbx {
         }
 
         KOKKOS_INLINE_FUNCTION NumType length() const {
-            return ::Kokkos::Experimental::sqrt(length_sqr());
+            return ::Kokkos::sqrt(length_sqr());
         }
 
         KOKKOS_INLINE_FUNCTION NumType dot(const vector3<NumType>& v) const {
@@ -177,8 +177,8 @@ namespace kokkostbx {
 
         // rotate a point about a unit vector3 axis
         KOKKOS_INLINE_FUNCTION vector3<NumType> rotate_around_axis(const vector3<NumType>& axis, NumType angle) const {
-            NumType sinphi = ::Kokkos::Experimental::sin(angle);
-            NumType cosphi = ::Kokkos::Experimental::cos(angle);
+            NumType sinphi = ::Kokkos::sin(angle);
+            NumType cosphi = ::Kokkos::cos(angle);
             NumType dot_factor = axis.dot(*this) * (1.0-cosphi);
 
             vector3<NumType> vector_rot = axis.cross(*this) * sinphi;

diff --git a/kokkostbx/SConscript b/kokkostbx/SConscript
@@ -22,7 +22,7 @@ if env_etc.enable_kokkos:
   if os.getenv('KOKKOS_ARCH') is None:
     os.environ['KOKKOS_ARCH'] = "HSW"
   if use_cuda and os.getenv('KOKKOS_CUDA_OPTIONS') is None:
-    os.environ['KOKKOS_CUDA_OPTIONS'] = "enable_lambda,force_uvm"
+    os.environ['KOKKOS_CUDA_OPTIONS'] = "enable_lambda"
   os.environ['CXXFLAGS'] = '-O3 -fPIC -DCUDAREAL=double'
 
   library_flags = "-Llib"
@@ -37,10 +37,6 @@ if env_etc.enable_kokkos:
     linked_libraries += " -lcudart -lcuda"
   os.environ['LDLIBS'] = linked_libraries
 
-  cxx_standard = '14'
-  if use_sycl:
-    cxx_standard = '17'
-
   original_cxx = None
   kokkos_lib = 'libkokkos.a'
   kokkos_cxxflags = None
@@ -119,7 +115,7 @@ if env_etc.enable_kokkos:
     returncode = subprocess.call([
         'cmake',
         os.environ['KOKKOS_PATH'],
-        '-DCMAKE_CXX_STANDARD={}'.format(cxx_standard),
+        '-DCMAKE_CXX_STANDARD={}'.format('17'),
         '-DCMAKE_INSTALL_PREFIX={}'.format(libtbx.env.under_build('.')),
         '-DCMAKE_INSTALL_LIBDIR=lib',
         '-DBUILD_SHARED_LIBS={}'.format(OnOff[True]),
@@ -138,9 +134,9 @@ if env_etc.enable_kokkos:
         '-DKokkos_ENABLE_SERIAL=ON',
         '-DKokkos_ENABLE_OPENMP={}'.format(OnOff[use_openmp]),
         '-DKokkos_ENABLE_CUDA={}'.format(OnOff[use_cuda]),
-        '-DKokkos_ENABLE_CUDA_UVM={}'.format(OnOff[use_cuda]),
         '-DKokkos_ENABLE_HIP={}'.format(OnOff[use_hip]),
-        '-DKokkos_ENABLE_SYCL={}'.format(OnOff[use_sycl])
+        '-DKokkos_ENABLE_SYCL={}'.format(OnOff[use_sycl]),
+        '-DKokkos_ENABLE_IMPL_MDSPAN=ON'
       ],
       cwd=kokkos_build_dir)
 

diff --git a/kokkostbx/kokkos_types.h b/kokkostbx/kokkos_types.h
@@ -3,10 +3,10 @@
 #include <Kokkos_Core.hpp>
 
 #ifdef KOKKOS_ENABLE_CUDA
-    #define MemSpace Kokkos::CudaUVMSpace
+    #define MemSpace Kokkos::CudaSpace
 #endif
 #ifdef KOKKOS_ENABLE_HIP
-    #define MemSpace Kokkos::Experimental::HIPSpace
+    #define MemSpace Kokkos::HIPSpace
 #endif
 #ifdef KOKKOS_ENABLE_OPENMPTARGET
     #define MemSpace Kokkos::OpenMPTargetSpace
@@ -19,8 +19,11 @@
 using ExecSpace = MemSpace::execution_space;
 using range_policy = Kokkos::RangePolicy<ExecSpace>;
 
-template <typename T>
-using view_1d_t = Kokkos::View<T*, MemSpace>;
+template <typename T> using view_1d_t = Kokkos::View<T*, MemSpace>;
+template <typename T> using view_4d_t = Kokkos::View<T****, MemSpace>;
+template <typename T> using view_5d_t = Kokkos::View<T*****, MemSpace>;
+template <typename T> using view_6d_t = Kokkos::View<T******, MemSpace>;
+template <typename T> using view_6d6_t = Kokkos::View<T******[6], MemSpace>;
 
 using vector_bool_t = view_1d_t<bool>;
 using vector_double_t = view_1d_t<double>;

diff --git a/kokkostbx/kokkos_utils.cpp b/kokkostbx/kokkos_utils.cpp
@@ -18,4 +18,20 @@ void transfer_double2kokkos(vector_cudareal_t& dst, const double* src, const siz
     }
 }
 
+void transfer_vector2kokkos(view_1d_t<bool>& dst, const std::vector<bool>& src) {
+    if (true) {
+        // printf("== Transfer %s from %p\n", dst.label().c_str(), (void*) dst.data());
+        // printf(" - size src|dst: %d|%d\n", src.size(), dst.span() );
+    }
+    if (dst.span() < src.size()) {
+        resize(dst, src.size());
+        // printf(" - size changed, new size: %d\n", dst.span() );
+    }
+    auto host_view = Kokkos::create_mirror_view(dst);
+    for (int i = 0; i < src.size(); ++i) {
+        host_view(i) = src[i];
+    }
+    Kokkos::deep_copy(dst, host_view);
+}
+
 }  // namespace kokkostbx
diff --git a/kokkostbx/kokkos_utils.h b/kokkostbx/kokkos_utils.h
@@ -78,6 +78,8 @@ void transfer_kokkos2shared(af::shared<T>& dst, const view_1d_t<T>& src) {
     transfer_kokkos2X(dst, src);
 }
 
+void transfer_vector2kokkos(view_1d_t<bool>& dst, const std::vector<bool>& src);
+
 template <typename T>
 void transfer_vector2kokkos(view_1d_t<T>& dst, const std::vector<T>& src) {
     if (true) {
@@ -88,11 +90,9 @@ void transfer_vector2kokkos(view_1d_t<T>& dst, const std::vector<T>& src) {
         resize(dst, src.size());
         // printf(" - size changed, new size: %d\n", dst.span() );
     }
-    auto host_view = Kokkos::create_mirror_view(dst);
-    for (int i = 0; i < src.size(); ++i) {
-        host_view(i) = src[i];
-    }
-    Kokkos::deep_copy(dst, host_view);
+    auto host_view = Kokkos::View<const T*, Kokkos::HostSpace>(src.data(), src.size());
+    auto dst_subview = Kokkos::subview(dst, std::pair<int, int>(0, src.size()));
+    Kokkos::deep_copy(dst_subview, host_view);
 }
 
 template <typename T>

diff --git a/kokkostbx/kokkos_vector.h b/kokkostbx/kokkos_vector.h
@@ -207,8 +207,9 @@ struct vector_base {
     }
 
     KOKKOS_FUNCTION void operator/=(const NumType& v) {
+        const NumType v_r = 1 / v;
         for (size_t i = 0; i < size; ++i) {
-            data[i] /= v;
+            data[i] *= v_r;
         }
     }
 
@@ -261,14 +262,15 @@ struct vector_base {
 
     KOKKOS_FUNCTION NumType length() const {
         // return sqrt_func(length_sqr());
-        return ::Kokkos::Experimental::sqrt(length_sqr());
+        return ::Kokkos::sqrt(length_sqr());
     }
 
     KOKKOS_FUNCTION void normalize() {
         NumType l = length();
         if (l > 0) {
+            NumType l_r = 1 / l;
             for (size_t i = 0; i < size; ++i) {
-                data[i] /= l;
+                data[i] *= l_r;
             }
         }
     }
@@ -277,8 +279,9 @@ struct vector_base {
         NumType l = length();
         Derived unit_vector{};
         if (l > 0) {
+            NumType l_r = 1 / l;
             for (size_t i = 0; i < size; ++i) {
-                unit_vector[i] = data[i] / l;
+                unit_vector[i] = data[i] * l_r;
             }
         }
         return unit_vector;

diff --git a/kokkostbx/kokkos_vector3.h b/kokkostbx/kokkos_vector3.h
@@ -10,8 +10,8 @@
 
 // #ifdef KOKKOS_CORE_HPP
 //     template <typename T> KOKKOS_FUNCTION T sin_func(T x) { return
-//     ::Kokkos::Experimental::sin(x); } template <typename T> KOKKOS_FUNCTION T cos_func(T x) {
-//     return ::Kokkos::Experimental::cos(x); }
+//     ::Kokkos::sin(x); } template <typename T> KOKKOS_FUNCTION T cos_func(T x) {
+//     return ::Kokkos::cos(x); }
 // #else
 //     #include <cmath>
 //     template <typename T> KOKKOS_FUNCTION T sin_func(T x) { return sin(x); }
@@ -27,26 +27,26 @@ struct vector3 : public vector<NumType, 3> {
     using vector_base = kokkostbx::vector<NumType, 3>;
 
     vector3() = default;
-    KOKKOS_FUNCTION vector3(NumType val) : vector_base(val){};
-    KOKKOS_FUNCTION vector3(NumType arr[]) : vector_base(arr){};
-    KOKKOS_FUNCTION vector3(const vector_base& vec) : vector_base(vec){};
+    KOKKOS_INLINE_FUNCTION vector3(NumType val) : vector_base(val){};
+    KOKKOS_INLINE_FUNCTION vector3(NumType arr[]) : vector_base(arr){};
+    KOKKOS_INLINE_FUNCTION vector3(const vector_base& vec) : vector_base(vec){};
 
-    KOKKOS_FUNCTION vector3(NumType x, NumType y, NumType z) : vector_base() {
+    KOKKOS_INLINE_FUNCTION vector3(NumType x, NumType y, NumType z) : vector_base() {
         vector_base::data[0] = x;
         vector_base::data[1] = y;
         vector_base::data[2] = z;
     }
 
     // decided against using properties, as this would increase the size of the class
-    KOKKOS_FUNCTION NumType& x_val() { return vector_base::data[0]; }
-    KOKKOS_FUNCTION NumType& y_val() { return vector_base::data[1]; }
-    KOKKOS_FUNCTION NumType& z_val() { return vector_base::data[2]; }
+    KOKKOS_INLINE_FUNCTION NumType& x_val() { return vector_base::data[0]; }
+    KOKKOS_INLINE_FUNCTION NumType& y_val() { return vector_base::data[1]; }
+    KOKKOS_INLINE_FUNCTION NumType& z_val() { return vector_base::data[2]; }
 
-    KOKKOS_FUNCTION NumType x_val() const { return vector_base::data[0]; }
-    KOKKOS_FUNCTION NumType y_val() const { return vector_base::data[1]; }
-    KOKKOS_FUNCTION NumType z_val() const { return vector_base::data[2]; }
+    KOKKOS_INLINE_FUNCTION NumType x_val() const { return vector_base::data[0]; }
+    KOKKOS_INLINE_FUNCTION NumType y_val() const { return vector_base::data[1]; }
+    KOKKOS_INLINE_FUNCTION NumType z_val() const { return vector_base::data[2]; }
 
-    KOKKOS_FUNCTION vector3<NumType> cross(const vector3<NumType>& v) const {
+    KOKKOS_INLINE_FUNCTION vector3<NumType> cross(const vector3<NumType>& v) const {
         vector3<NumType> cross_vector{};
         cross_vector.x_val() = y_val() * v.z_val() - z_val() * v.y_val();
         cross_vector.y_val() = z_val() * v.x_val() - x_val() * v.z_val();
@@ -56,14 +56,14 @@ struct vector3 : public vector<NumType, 3> {
     }
 
     // rotate a point around a unit vector3 axis
-    KOKKOS_FUNCTION vector3<NumType> rotate_around_axis(const vector3<NumType>& axis, NumType angle)
+    KOKKOS_INLINE_FUNCTION vector3<NumType> rotate_around_axis(const vector3<NumType>& axis, NumType angle)
         const {
         // NumType sinphi = sin_func(angle);
         // NumType cosphi = cos_func(angle);
-        NumType sinphi = ::Kokkos::Experimental::sin(angle);
-        NumType cosphi = ::Kokkos::Experimental::cos(angle);
+        const NumType sinphi = ::Kokkos::sin(angle);
+        const NumType cosphi = ::Kokkos::cos(angle);
 
-        NumType dot_factor = axis.dot(*this) * (1.0 - cosphi);
+        const NumType dot_factor = axis.dot(*this) * (1.0 - cosphi);
 
         vector3<NumType> vector_rot = axis.cross(*this) * sinphi;
         vector_rot += axis * dot_factor;

diff --git a/libtbx/auto_build/bootstrap.py b/libtbx/auto_build/bootstrap.py
@@ -1006,17 +1006,17 @@ class xia2_module(SourceModule):
 
 class kokkos_module(SourceModule):
   module = 'kokkos'
-  anonymous = ['git', '-b 3.7.01',
+  anonymous = ['git', '-b 4.2.00',
                '[email protected]:kokkos/kokkos.git',
                'https://github.com/kokkos/kokkos.git',
-               'https://github.com/kokkos/kokkos/archive/refs/tags/3.7.01.zip']
+               'https://github.com/kokkos/kokkos/archive/refs/tags/4.2.00.zip']
 
 class kokkos_kernels_module(SourceModule):
   module = 'kokkos-kernels'
-  anonymous = ['git', '-b 3.7.01',
+  anonymous = ['git', '-b 4.2.00',
                '[email protected]:kokkos/kokkos-kernels.git',
                'https://github.com/kokkos/kokkos-kernels.git',
-               'https://github.com/kokkos/kokkos-kernels/archive/refs/tags/3.7.01.zip']
+               'https://github.com/kokkos/kokkos-kernels/archive/refs/tags/4.2.00.zip']
 
 # Duke repositories
 class probe_module(SourceModule):

diff --git a/simtbx/command_line/complete_an_F.py b/simtbx/command_line/complete_an_F.py
@@ -0,0 +1,48 @@
+from __future__ import division, print_function
+from argparse import ArgumentParser
+parser = ArgumentParser()
+parser.add_argument("mtzin", help="input mtz file", type=str)
+parser.add_argument("mtzout", help="output mtz file", type=str)
+args = parser.parse_args()
+
+# LIBTBX_SET_DISPATCHER_NAME diffBragg.completeF
+
+import numpy as np
+from iotbx.reflection_file_reader import any_reflection_file
+from dials.array_family import flex
+from scipy.interpolate import interp1d
+from cctbx import miller
+
+F = any_reflection_file(args.mtzin).as_miller_arrays()[0]
+F = F.as_amplitude_array()
+if not F.is_xray_amplitude_array():
+    F = F.set_observation_type_xray_amplitude()
+
+print("Bin-ID    Res-range    Completeness    #ASU-indices")
+F.show_completeness()
+d_max,d_min = F.resolution_range()
+print("d_min, d_max (Angstrom): ", d_min, d_max)
+mset_full = F.build_miller_set(False, d_min=d_min)
+mset_full_d = {h: d for h,d in zip(mset_full.d_spacings().indices(), mset_full.d_spacings().data())}
+Fmap = {h:val for h,val in zip(F.indices(), F.data())}
+xvals = np.array(F.d_spacings().data())
+yvals = np.array(F.data())
+fill_vals = yvals[np.argmin(xvals)], yvals[np.argmax(xvals)]
+I = interp1d(xvals, yvals, fill_value=fill_vals, bounds_error=False)
+data = []
+for h in mset_full.indices():
+    if h not in Fmap:
+        d_h = mset_full_d[h]
+        amp = I(d_h)
+    else:
+        amp = Fmap[h]
+    data.append(amp)
+
+complete_amps = flex.double(data)
+complete_inds = mset_full.indices()
+ma = miller.array(mset_full, complete_amps)
+if not ma.is_xray_amplitude_array():
+    ma = ma.set_observation_type_xray_amplitude()
+ma = ma.as_anomalous_array()
+assert ma.anomalous_flag()
+ma.as_mtz_dataset(column_root_label="F").mtz_object().write(args.mtzout)
-Original file line number
+Diff line change
@@ Expand Up / @@ -95,6 +95,7 @@ steps: @@
           libtbx.configure lunus
           make
         fi
+        rm -fr modules/kokkos/.git/*
       displayName: Configure and Build
     # test
@@ Expand Down @@