try fixing early thread exit in ndeg twisted mass

Author: jcosborn

include/kernels/dslash_ndeg_twisted_mass_preconditioned.cuh

@@ -73,6 +73,7 @@ namespace quda
       int src_idx = src_flavor / 2;
       int flavor = src_flavor % 2;
 
+      bool active2 = mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
       int thread_dim; // which dimension is thread working on (fused kernel only)
       auto coord = getCoords<QUDA_4D_PC, mykernel_type>(arg, idx, flavor, parity, thread_dim);
 
@@ -86,16 +87,15 @@ namespace quda
           active = false;
         }
       }
-      active &= mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
 
       if (!allthreads || active) {
 	if (!dagger || Arg::asymmetric) // defined in dslash_wilson.cuh
-	  applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
+	  applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active2, src_idx);
 	else // defined in dslash_twisted_mass_preconditioned
-	  applyWilsonTM<dagger, 2, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
+	  applyWilsonTM<dagger, 2, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active2, src_idx);
       }
 
-      if (active) {
+      if (!allthreads || active) {
         if (xpay && mykernel_type == INTERIOR_KERNEL && !arg.dd_x.isZero(coord)) {
           if constexpr (!dagger || Arg::asymmetric) { // apply inverse twist which is undone below
             // use consistent load order across s to ensure better cache locality
@@ -109,7 +109,7 @@ namespace quda
             Vector x = arg.x[src_idx](my_flavor_idx, my_spinor_parity);
             out += x; // just directly add since twist already applied in the dslash
           }
-        } else if (mykernel_type != INTERIOR_KERNEL) {
+        } else if (mykernel_type != INTERIOR_KERNEL && active2) {
           // if we're not the interior kernel, then we must sum the partial
           Vector x = arg.out[src_idx](my_flavor_idx, my_spinor_parity);
           out += x;
@@ -118,20 +118,24 @@ namespace quda
 
       if constexpr (!dagger || Arg::asymmetric) { // apply A^{-1} to D*in
         SharedMemoryCache<Vector> cache {*this};
-        if (isComplete<mykernel_type>(arg, coord) && active) {
-          // to apply the preconditioner we need to put "out" in shared memory so the other flavor can access it
-          cache.save(out);
-        }
+	if (!allthreads || active) {
+	  if (isComplete<mykernel_type>(arg, coord) && active2) {
+	    // to apply the preconditioner we need to put "out" in shared memory so the other flavor can access it
+	    cache.save(out);
+	  }
+	}
         cache.sync(); // safe to sync here since other threads will exit if allowed, or all be here
-        if (isComplete<mykernel_type>(arg, coord) && active) {
-          if (flavor == 0)
-            out = arg.a * (out + arg.b * out.igamma(4) + arg.c * cache.load_y(target::thread_idx().y + 1));
-          else
-            out = arg.a * (out - arg.b * out.igamma(4) + arg.c * cache.load_y(target::thread_idx().y - 1));
-        }
+	if (!allthreads || active) {
+	  if (isComplete<mykernel_type>(arg, coord) && active2) {
+	    if (flavor == 0)
+	      out = arg.a * (out + arg.b * out.igamma(4) + arg.c * cache.load_y(target::thread_idx().y + 1));
+	    else
+	      out = arg.a * (out - arg.b * out.igamma(4) + arg.c * cache.load_y(target::thread_idx().y - 1));
+	  }
+	}
       }
 
-      if (active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
+      if (active && active2) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
     }
   };