Feature/DD

CMakeLists.txt

@@ -151,8 +151,8 @@ option(QUDA_DIRAC_CLOVER_HASENBUSCH "build clover Hasenbusch twist operators" ${
 option(QUDA_DIRAC_NDEG_TWISTED_MASS "build non-degenerate twisted mass Dirac operators" ${QUDA_DIRAC_DEFAULT})
 option(QUDA_DIRAC_NDEG_TWISTED_CLOVER "build non-degenerate twisted clover Dirac operators" ${QUDA_DIRAC_DEFAULT})
 option(QUDA_DIRAC_LAPLACE "build laplace operator" ${QUDA_DIRAC_DEFAULT})
-
 option(QUDA_DIRAC_DISTANCE_PRECONDITIONING "build code for distance preconditioned Wilson/clover Dirac operators" OFF)
+set(QUDA_DOMAIN_DECOMPOSITION "0" CACHE STRING "which domain decomposition to instantiate in QUDA (1-bit number - RedBlack)")
 
 option(QUDA_COVDEV "build code for covariant derivative" OFF)
 

ci/docker/Dockerfile.build

@@ -39,6 +39,7 @@ RUN  QUDA_TEST_GRID_SIZE="1 1 1 2" cmake -S /quda/src \
     -DQUDA_MULTIGRID_NVEC_LIST=6 \
     -DQUDA_MDW_FUSED_LS_LIST=4 \
     -DQUDA_MPI=ON  \
+    -DQUDA_DSLASH_DISTANCE=1 \
     -DQUDA_DIRAC_DEFAULT_OFF=ON \
     -DQUDA_DIRAC_WILSON=ON \
     -DQUDA_DIRAC_CLOVER=ON \

include/color_spinor_field.h

@@ -149,6 +149,7 @@ namespace quda
     int composite_dim = 0; // e.g., number of eigenvectors in the set
     bool is_component = false;
     int component_id = 0; // eigenvector index
+    DDParam dd {};
 
     /**
        If using CUDA native fields, this function will ensure that the
@@ -356,6 +357,9 @@ namespace quda
     //
     CompositeColorSpinorField components;
 
+    /** Domain decomposition options */
+    DDParam dd {};
+
     /**
        Compute the required extended ghost zone sizes and offsets
        @param[in] nFace The depth of the halo
@@ -438,6 +442,31 @@ namespace quda
      */
     void copy(const ColorSpinorField &src);
 
+    /**
+       @brief Project the field to a domain determined by DDParam
+     */
+    void projectDD();
+
+    /**
+       @brief Returns DDParam (const version)
+     */
+    const DDParam& DD() const { return dd; }
+
+    /**
+       @brief Returns DDParam (non const version)
+     */
+    DDParam& DD() { return dd; }
+
+    /**
+       @brief Sets DDParam from a given DDParam
+     */
+    void DD(const DDParam &in) { dd = in; }
+
+    /**
+       @brief Sets DDParam from a given list of options (DD flags)
+     */
+    template <typename... Args> void DD(const quda::DD &flag, const Args &...args) { dd.set(flag, args...); }
+
     /**
        @brief Zero all elements of this field
      */
@@ -977,6 +1006,8 @@ namespace quda
                               void *Dst = nullptr, const void *Src = nullptr);
 
   void genericSource(ColorSpinorField &a, QudaSourceType sourceType, int x, int s, int c);
+
+  void genericProjectDD(ColorSpinorField &a);
   int genericCompare(const ColorSpinorField &a, const ColorSpinorField &b, int tol);
 
   /**

include/declare_enum.h

@@ -0,0 +1,60 @@
+/*
+ * A macro that declares an `enum class` as well as a `to_string` function for the enums.
+ * The enum has also a default value `size` that measures the size of the enum.
+ *
+ * Credit: https://stackoverflow.com/a/71375077/12084612
+ * -------
+ * License: CC BY-SA 4.0
+ * --------
+ * Usage:
+ * ------
+ *
+ * DECLARE_ENUM(WeekEnum, Mon, Tue, Wed, Thu, Fri, Sat, Sun,);
+ *
+ * int main()
+ * {
+ *     WeekEnum weekDay = WeekEnum::Wed;
+ *     std::cout << to_string(weekDay) << std::endl; // prints Wed
+ *     std::cout << to_string(WeekEnum::Sat) << std::endl; // prints Sat
+ *     std::cout << to_string((int) WeekEnum::size) << std::endl; // prints 7
+ *     return 0;
+ * }
+ *
+ */
+
+#pragma once
+#include <iostream>
+#include <string>
+#include <vector>
+#include <algorithm>
+#include <unordered_map>
+
+// Add the definition of this method into a cpp file. (only the declaration in the header)
+static inline const std::vector<std::string> get_enum_names(const std::string &en_key, const std::string &en_str)
+{
+  static std::unordered_map<std::string, std::vector<std::string>> en_names_map;
+  const auto it = en_names_map.find(en_key);
+  if (it != en_names_map.end()) return it->second;
+
+  constexpr auto delim(',');
+  std::vector<std::string> en_names;
+  std::size_t start {};
+  auto end = en_str.find(delim);
+  while (end != std::string::npos) {
+    while (en_str[start] == ' ') ++start;
+    en_names.push_back(en_str.substr(start, end - start));
+    start = end + 1;
+    end = en_str.find(delim, start);
+  }
+  while (en_str[start] == ' ') ++start;
+  en_names.push_back(en_str.substr(start));
+  return en_names_map.emplace(en_key, std::move(en_names)).first->second;
+}
+
+#define DECLARE_ENUM(ENUM_NAME, ...)                                                                                   \
+  enum class ENUM_NAME : unsigned int { __VA_ARGS__ size };                                                            \
+  inline std::string to_string(ENUM_NAME en)                                                                           \
+  {                                                                                                                    \
+    const auto names = get_enum_names(#ENUM_NAME, #__VA_ARGS__);                                                       \
+    return names[static_cast<std::size_t>(en)];                                                                        \
+  }

include/domain_decomposition.h

@@ -0,0 +1,153 @@
+#pragma once
+
+#include "declare_enum.h"
+
+namespace quda
+{
+
+  // using namespace quda;
+
+  DECLARE_ENUM(DD, // name of the enum class
+
+               reset, // No domain decomposition. It sets all flags to zero.
+
+               red_black_type,   // Flags used by red_black
+               red_active,       // if red blocks are active
+               black_active,     // if black blocks are active
+               no_block_hopping, // if hopping between red and black is allowed
+  );
+
+  // Params for domain decompation
+  struct DDParam {
+
+    QudaDDType type = QUDA_DD_NO;
+    array<bool, static_cast<int>(DD::size)> flags = {}; // the default value of all flags is 0
+    array<int, QUDA_MAX_DIM> block_dim = {};            // the size of the block per direction
+
+    // Default constructor
+    DDParam() = default;
+
+    // returns false if in use
+    constexpr bool operator!() const { return type == QUDA_DD_NO; }
+
+    // returns value of given flag
+    constexpr bool is(const DD &flag) const { return flags[(int)flag]; }
+
+    // sets given flag to true
+    constexpr void set(const DD &flag)
+    {
+      flags[(int)flag] = true;
+
+      if ((int)flag == (int)DD::reset) {
+#pragma unroll
+        for (auto i = 0u; i < (int)DD::size; i++) flags[i] = 0;
+        type = QUDA_DD_NO;
+      } else if ((int)flag >= (int)DD::red_black_type) {
+        type = QUDA_DD_RED_BLACK;
+      }
+    }
+
+    template <typename... Args> constexpr void set(const DD &flag, const Args &...args)
+    {
+      set(flag);
+      set(args...);
+    }
+
+    // Pretty print the args struct
+    void print() const
+    {
+      if (not *this) {
+        printfQuda("DD not in use\n");
+        return;
+      }
+      printfQuda("Printing DDParam\n");
+      for (int i = 0; i < (int)DD::size; i++)
+        printfQuda("flags[DD::%s] = %s\n", to_string((DD)i).c_str(), flags[i] ? "true" : "false");
+      for (int i = 0; i < QUDA_MAX_DIM; i++) printfQuda("block_dim[%d] = %d\n", i, static_cast<int>(block_dim[i]));
+    }
+
+    // Checks if this matches to given DDParam
+    template <typename F> inline bool check(const F &field, bool verbose = false) const
+    {
+      if (not *this) return true;
+
+      if (type == QUDA_DD_RED_BLACK) {
+        for (int i = 0; i < field.Ndim(); i++) {
+          if (block_dim[i] < 0) {
+            if (verbose) printfQuda("block_dim[%d] = %d is negative\n", i, block_dim[i]);
+            return false;
+          }
+          if (block_dim[i] > 0) {
+            int globalDim = comm_dim(i) * field.full_dim(i);
+            if (globalDim % block_dim[i] != 0) {
+              if (verbose) printfQuda("block_dim[%d] = %d does not divide %d \n", i, block_dim[i], globalDim);
+              return false;
+            }
+            if ((globalDim / block_dim[i]) % 2 != 0) {
+              if (verbose)
+                printfQuda("block_dim[%d] = %d does not divide %d **evenly** \n", i, block_dim[i], globalDim);
+              return false;
+            }
+          }
+        }
+        if (block_dim[0] % 2) {
+          if (verbose) printfQuda("block_dim[0] = %d must be even \n", block_dim[0]);
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    // Checks if this matches to given DDParam
+    inline bool match(const DDParam &dd, bool verbose = false) const
+    {
+      // if one of the two is not in use we return true, i.e. one of the two is a full field
+      if (not *this or not dd) return true;
+
+      // false if type does not match
+      if (type != dd.type) {
+        if (verbose) printfQuda("DD type do not match (%d != %d)\n", type, dd.type);
+        return false;
+      }
+
+      if (type == QUDA_DD_RED_BLACK) {
+        for (int i = 0; i < QUDA_MAX_DIM; i++)
+          if (block_dim[i] != dd.block_dim[i]) {
+            if (verbose) printfQuda("block_dim[%d] = %d != %d \n", i, block_dim[i], dd.block_dim[i]);
+            return false;
+          }
+        if (is(DD::no_block_hopping) != dd.is(DD::no_block_hopping)) {
+          if (verbose) printfQuda("no_block_hopping do not match.\n");
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    // Checks if this is equal to given DDParam
+    inline bool operator==(const DDParam &dd) const
+    {
+      // if both are not in use we return true
+      if (not *this and not dd) return true;
+
+      // false if type does not match
+      if (type != dd.type) return false;
+
+      // checking all flags matches (note this should be actually type-wise)
+      for (int i = 0; i < (int)DD::size; i++)
+        if (flags[i] != dd.flags[i]) return false;
+
+      // checking block_dim matches when needed
+      if (type == QUDA_DD_RED_BLACK)
+        for (int i = 0; i < QUDA_MAX_DIM; i++)
+          if (block_dim[i] != dd.block_dim[i]) return false;
+
+      return true;
+    }
+
+    inline bool operator!=(const DDParam &dd) const { return !(*this == dd); }
+  };
+
+} // namespace quda