Feature/prefetch2

CMakeLists.txt

@@ -239,7 +239,7 @@ option(QUDA_CTEST_SEP_DSLASH_POLICIES "Test Dslash policies separately in ctest
 option(QUDA_CTEST_DISABLE_BENCHMARKS "Disable benchmark test" ON)
 
 option(QUDA_FAST_COMPILE_REDUCE "enable fast compilation in blas and reduction kernels (single warp per reduction)" OFF)
-option(QUDA_FAST_COMPILE_DSLASH "enable fast compilation in dslash kernels (~20% perf impact)" OFF)
+option(QUDA_FAST_COMPILE_DSLASH "enable fast compilation in coarse grid dslash kernels (significant perf impact)" OFF)
 
 option(QUDA_OPENMP "enable OpenMP" OFF)
 set(QUDA_CXX_STANDARD

ci/docker/Dockerfile.build

@@ -8,7 +8,9 @@ RUN apt-get update -qq && apt-get install -qq -y --no-install-recommends \
     build-essential \
     cmake \
     wget \
-    ninja-build && \
+    ninja-build \
+    git \
+    ca-certificates && \
     rm -rf /var/lib/apt/lists/*
 
 ARG MPICH_VERSION=3.3.2

include/color_spinor_field_order.h

@@ -241,9 +241,9 @@ namespace quda
         constexpr int M = nSpinBlock * nColor * nVec;
 #pragma unroll
         for (int i = 0; i < M; i++) {
-          vec_t tmp
-            = vector_load<vec_t>(reinterpret_cast<const vec_t *>(in + parity * offset_cb), x_cb * N + chi * M + i);
-          memcpy(&out[i], &tmp, sizeof(vec_t));
+          auto tmp
+            = vector_load<Float, 2>(reinterpret_cast<const vec_t *>(in + parity * offset_cb), x_cb * N + chi * M + i);
+          memcpy(&out[i], &tmp, sizeof(tmp));
         }
       }
     };
@@ -1010,11 +1010,14 @@ namespace quda
       {
         for (int dim = 0; dim < 4; dim++) {
           for (int dir = 0; dir < 2; dir++) {
-            ghost[2 * dim + dir] = comm_dim_partitioned(dim) ? static_cast<Float *>(ghost_[2 * dim + dir]) : nullptr;
-            ghost_norm[2 * dim + dir] = !comm_dim_partitioned(dim) ?
-              nullptr :
-              reinterpret_cast<norm_type *>(static_cast<char *>(ghost_[2 * dim + dir])
-                                            + nParity * length_ghost * faceVolumeCB[dim] * sizeof(Float));
+            if (comm_dim_partitioned(dim) && ghost_[2 * dim + dir]) {
+              ghost[2 * dim + dir] = static_cast<Float *>(ghost_[2 * dim + dir]);
+              ghost_norm[2 * dim + dir] = reinterpret_cast<norm_type *>(
+                static_cast<char *>(ghost_[2 * dim + dir]) + nParity * length_ghost * faceVolumeCB[dim] * sizeof(Float));
+            } else {
+              ghost[2 * dim + dir] = nullptr;
+              ghost_norm[2 * dim + dir] = nullptr;
+            }
           }
         }
       }
@@ -1023,7 +1026,7 @@ namespace quda
       {
         real v[length_ghost];
         norm_type nrm
-          = isFixed<Float>::value ? vector_load<float>(ghost_norm[2 * dim + dir], parity * faceVolumeCB[dim] + x) : 0.0;
+          = isFixed<Float>::value ? vector_load<float, 1>(ghost_norm[2 * dim + dir], parity * faceVolumeCB[dim] + x)[0] : 0.0;
 
 #pragma unroll
         for (int i = 0; i < M; i++) {
@@ -1123,16 +1126,9 @@ namespace quda
       using real = typename mapper<Float>::type;
       using complex = complex<real>;
       using AllocInt = typename AllocType<huge_alloc>::type;
-      using norm_type = float;
+      using norm_t = float;
       Float *field = nullptr;
-      //#define LEGACY_ACCESSOR_NORM // legacy code where norm pointer and offset are stored instead of computed
-#ifdef LEGACY_ACCESSOR_NORM
-      norm_type *norm = nullptr;
-#endif
       AllocInt offset = 0; // offset can be 32-bit or 64-bit
-#ifdef LEGACY_ACCESSOR_NORM
-      AllocInt norm_offset = 0;
-#endif
       int volumeCB = 0;
 
       FloatNOrder() = default;
@@ -1141,14 +1137,7 @@ namespace quda
       FloatNOrder(const ColorSpinorField &a, int nFace = 1, Float *buffer = 0, Float **ghost_ = 0) :
         GhostNOrder(a, nFace, ghost_),
         field(buffer ? buffer : a.data<Float *>()),
-#ifdef LEGACY_ACCESSOR_NORM
-        norm(buffer ? reinterpret_cast<norm_type *>(reinterpret_cast<char *>(buffer) + a.NormOffset()) :
-                      const_cast<norm_type *>(reinterpret_cast<const norm_type *>(a.Norm()))),
-#endif
         offset(a.Bytes() / (2 * sizeof(Float))),
-#ifdef LEGACY_ACCESSOR_NORM
-        norm_offset(a.Bytes() / (2 * sizeof(norm_type))),
-#endif
         volumeCB(a.VolumeCB())
       {
       }
@@ -1157,54 +1146,59 @@ namespace quda
       __device__ __host__ inline void load(complex out[length / 2], int x, int parity = 0) const
       {
         real v[length];
-#ifndef LEGACY_ACCESSOR_NORM
-        auto norm_offset = offset / (sizeof(Float) < sizeof(float) ? sizeof(norm_type) / sizeof(Float) : 1);
-        auto norm = reinterpret_cast<float *>(field + volumeCB * (2 * Nc * Ns));
-#endif
-        norm_type nrm = isFixed<Float>::value ? vector_load<float>(norm, x + parity * norm_offset) : 0.0;
-
+        auto norm_offset = (volumeCB * 2 * Nc * Ns + parity * offset) * sizeof(Float) / sizeof(norm_t);
+        norm_t nrm = isFixed<Float>::value ? vector_load<norm_t, 1>(field, x + norm_offset)[0] : 0.0;
 #pragma unroll
         for (int i = 0; i < M; i++) {
           // first load from memory
-          auto vecTmp = vector_load<Float, N>(field + parity * offset, volumeCB * i + x);
+          auto vecTmp = vector_load<Float, N>(field, parity * offset, volumeCB * i + x);
           // now copy into output and scale
           copy_and_scale(v + i * N, vecTmp, nrm);
         }
 
         // now load any remainder
         if constexpr (Nrem > 0) {
-          auto vecTmp = vector_load<Float, Nrem>(field + parity * offset + volumeCB * M * N, x);
+          auto vecTmp = vector_load<Float, Nrem>(field, parity * offset + volumeCB * M * N, x);
           copy_and_scale(v + M * N, vecTmp, nrm);
         }
 
 #pragma unroll
         for (int i = 0; i < length / 2; i++) out[i] = complex(v[2 * i + 0], v[2 * i + 1]);
       }
 
+      __device__ __host__ inline void prefetch(int x, int parity = 0) const
+      {
+        auto norm_offset = (volumeCB * 2 * Nc * Ns + parity * offset) * sizeof(Float) / sizeof(norm_t);
+        if constexpr (isFixed<Float>::value) prefetch_cache_line(reinterpret_cast<norm_t *>(field) + (x + norm_offset));
+
+#pragma unroll
+        for (int i = 0; i < M; i++) prefetch_cache_line(field + (parity * offset + (volumeCB * i + x) * N));
+
+        // now load any remainder
+        if constexpr (Nrem > 0) prefetch_cache_line(field + (parity * offset + volumeCB * M * N + x * Nrem));
+      }
+
       __device__ __host__ inline void save(const complex in[length / 2], int x, int parity = 0) const
       {
         real v[length];
-#ifndef LEGACY_ACCESSOR_NORM
-        auto norm_offset = offset / (sizeof(Float) < sizeof(float) ? sizeof(norm_type) / sizeof(Float) : 1);
-        auto norm = reinterpret_cast<float *>(field + volumeCB * (2 * Nc * Ns));
-#endif
+        auto norm_offset = (volumeCB * 2 * Nc * Ns + parity * offset) * sizeof(Float) / sizeof(norm_t);
+
 #pragma unroll
         for (int i = 0; i < length / 2; i++) {
           v[2 * i + 0] = in[i].real();
           v[2 * i + 1] = in[i].imag();
         }
 
-        norm_type scale = 0.0;
-        norm_type scale_inv = 0.0;
+        norm_t scale = 0.0;
+        norm_t scale_inv = 0.0;
         if constexpr (isFixed<Float>::value) {
-          norm_type max_[length / 2];
+          norm_t max_[length / 2];
           // two-pass to increase ILP (assumes length divisible by two, e.g. complex-valued)
 #pragma unroll
-          for (int i = 0; i < length / 2; i++)
-            max_[i] = fmaxf(fabsf((norm_type)v[i]), fabsf((norm_type)v[i + length / 2]));
+          for (int i = 0; i < length / 2; i++) max_[i] = fmaxf(fabsf((norm_t)v[i]), fabsf((norm_t)v[i + length / 2]));
 #pragma unroll
           for (int i = 0; i < length / 2; i++) scale = fmaxf(max_[i], scale);
-          norm[x + parity * norm_offset] = scale * fixedInvMaxValue<Float>::value;
+          reinterpret_cast<norm_t *>(field)[x + norm_offset] = scale * fixedInvMaxValue<Float>::value;
           scale_inv = fdividef(fixedMaxValue<Float>::value, scale);
         }
 
@@ -1214,14 +1208,14 @@ namespace quda
           // first do scalar copy converting into storage type
           copy_and_scale<Float, real, N>(vecTmp, v + i * N, scale_inv);
           // second do vectorized copy into memory
-          vector_store(field + parity * offset, volumeCB * i + x, vecTmp);
+          vector_store(field, parity * offset, volumeCB * i + x, vecTmp);
         }
 
         if constexpr (Nrem > 0) {
           array<Float, Nrem> vecTmp;
           copy_and_scale<Float, real, Nrem>(vecTmp, v + M * N, scale_inv);
           // second do vectorized copy into memory
-          vector_store(field + parity * offset + volumeCB * M * N, x, vecTmp);
+          vector_store(field, parity * offset + volumeCB * M * N, x, vecTmp);
         }
       }
 

include/complex_quda.h

@@ -928,14 +928,14 @@ namespace quda
   template <typename real> __host__ __device__ inline complex<real> cmul(const complex<real> &x, const complex<real> &y)
   {
     complex<real> rtn = mul2({x.real(), x.real()}, y);
-    return fma2({x.imag(), x.imag()}, {-y.imag(), y.real()}, rtn);
+    return fma2({-x.imag(), x.imag()}, {y.imag(), y.real()}, rtn);
   }
 
   template <typename real>
   __host__ __device__ inline complex<real> cmac(const complex<real> &x, const complex<real> &y, const complex<real> &z)
   {
     complex<real> w = fma2({x.real(), x.real()}, y, z);
-    return fma2({x.imag(), x.imag()}, {-y.imag(), y.real()}, w);
+    return fma2({-x.imag(), x.imag()}, {y.imag(), y.real()}, w);
   }
 
   template <typename T1, typename T2, typename T3>

include/domain_decomposition.h

@@ -39,8 +39,7 @@ namespace quda
       flags[(int)flag] = true;
 
       if ((int)flag == (int)DD::reset) {
-#pragma unroll
-        for (auto i = 0u; i < (int)DD::size; i++) flags[i] = 0;
+        flags = {};
         type = QUDA_DD_NO;
       } else if ((int)flag >= (int)DD::red_black_type) {
         type = QUDA_DD_RED_BLACK;

include/dslash.h

@@ -8,6 +8,7 @@
 #include <tunable_nd.h>
 #include <instantiate.h>
 #include <instantiate_dslash.h>
+#include <tma_helper.hpp>
 
 namespace quda
 {
@@ -70,6 +71,18 @@ namespace quda
       char tile_str[16];
       i32toa(tile_str, Arg::n_src_tile);
       strcat(aux_base, tile_str);
+      if constexpr (dslash_double_store()) strcat(aux_base, ",double_store");
+      if constexpr (Arg::prefetch_distance > 0) {
+        strcat(aux_base, ",prefetch=");
+        i32toa(tile_str, Arg::prefetch_distance);
+        strcat(aux_base, tile_str);
+        if constexpr (dslash_prefetch_type() == PrefetchType::THREAD)
+          strcat(aux_base, ",prefetch=thread");
+        else if constexpr (dslash_prefetch_type() == PrefetchType::BULK)
+          strcat(aux_base, ",prefetch=bulk");
+        else if constexpr (dslash_prefetch_type() == PrefetchType::TENSOR)
+          strcat(aux_base, ",prefetch=tensor");
+      }
     }
 
     /**
@@ -130,7 +143,7 @@ namespace quda
       }
     }
 
-    inline void setParam(TuneParam &tp)
+    template <bool improved = false> inline void setParam(TuneParam &tp, const GaugeField &U, const GaugeField &L = {})
     {
       // Need to reset ghost pointers prior to every call since the
       // ghost buffer may have been changed during policy tuning.
@@ -173,6 +186,16 @@ namespace quda
           0;
         tp.grid.x += arg.exterior_blocks;
       }
+
+      if constexpr (dslash_prefetch_type() == PrefetchType::TENSOR && Arg::prefetch_distance > 0) {
+        Dslash::arg.U.tensor_desc = get_tensor_descriptor(U, tp.block.x);
+        Dslash::arg.Uback.tensor_desc = get_tensor_descriptor(U.shift(), tp.block.x);
+        if constexpr (improved) {
+          assert(!U.empty());
+          Dslash::arg.L.tensor_desc = get_tensor_descriptor(L, tp.block.x);
+          Dslash::arg.Lback.tensor_desc = get_tensor_descriptor(L.shift(), tp.block.x);
+        }
+      }
     }
 
     virtual int blockStep() const override { return (arg.shmem & 64) ? 8 : 16; }
@@ -219,6 +242,15 @@ namespace quda
       }
     }
 
+    virtual bool advanceBlockDim(TuneParam &param) const override
+    {
+      // if TMA is enabled we must keep parity separate in the block (2-d tuning)
+      if constexpr (dslash_prefetch_tma())
+        return TunableKernel2D_base<false>::advanceBlockDim(param);
+      else
+        return TunableKernel3D::advanceBlockDim(param);
+    }
+
     virtual bool advanceTuneParam(TuneParam &param) const override
     {
       return advanceAux(param) || advanceSharedBytes(param) || advanceBlockDim(param) || advanceSharedCarveOut(param)
@@ -268,6 +300,7 @@ namespace quda
     inline void launch(TuneParam &tp, const qudaStream_t &stream)
     {
       tp.set_max_shared_bytes = true;
+      if (dslash_prefetch_tma() && tp.block.z > 1) errorQuda("Z-dimension block size must be 1 when using TMA");
       launch_device<dslash_functor>(
         tp, stream, dslash_functor_arg<D, P, dagger, xpay, kernel_type, Arg>(arg, tp.block.x * tp.grid.x));
     }