Skip to content

Commit e95db1e

Browse files
jcosbornjcosborn
committed
fix inactive thread handling in ndeg twisted clover
1 parent 508ce26 commit e95db1e

3 files changed

Lines changed: 87 additions & 82 deletions

File tree

‎include/kernels/dslash_ndeg_twisted_clover.cuh‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ namespace quda
1414
static constexpr int length = (nSpin / (nSpin / 2)) * 2 * nColor * nColor * (nSpin / 2) * (nSpin / 2) / 2;
1515
typedef typename clover_mapper<Float, length, true>::type C;
1616
typedef typename mapper<Float>::type real;
17-
17+
1818
const C A; /** the clover field */
1919
real a; /** this is the Wilson-dslash scale factor */
2020
real b; /** this is the chiral twist factor */
@@ -59,14 +59,15 @@ namespace quda
5959
Note this routine only exists in xpay form.
6060
*/
6161
template <KernelType mykernel_type = kernel_type, bool allthreads = false>
62-
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool active = true)
62+
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool alive = true)
6363
{
6464
typedef typename mapper<typename Arg::Float>::type real;
6565
typedef ColorSpinor<real, Arg::nColor, 4> Vector;
6666
typedef ColorSpinor<real, Arg::nColor, 2> HalfVector;
6767

6868
int src_idx = src_flavor / 2;
6969
int flavor = src_flavor % 2;
70+
bool active = mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
7071
int thread_dim; // which dimension is thread working on (fused kernel only)
7172

7273
auto coord = getCoords<QUDA_4D_PC, mykernel_type>(arg, idx, flavor, parity, thread_dim);
@@ -75,27 +76,26 @@ namespace quda
7576
const int my_flavor_idx = coord.x_cb + flavor * arg.dc.volume_4d_cb;
7677
Vector out;
7778

78-
active &= mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
79-
if (arg.dd_out.isZero(coord)) {
80-
if (active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
81-
if constexpr (!allthreads) return;
82-
else active = false;
79+
if (!allthreads || alive) {
80+
if (arg.dd_out.isZero(coord)) {
81+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
82+
if constexpr (!allthreads) return;
83+
else alive = false;
84+
}
8385
}
8486

85-
if (!allthreads || active) {
87+
if (!allthreads || alive) {
8688
// defined in dslash_wilson.cuh
8789
applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
8890
}
8991

9092
if constexpr (mykernel_type == INTERIOR_KERNEL) {
9193
if (arg.dd_x.isZero(coord)) {
92-
if (!allthreads || active) {
93-
out = arg.a * out;
94-
}
94+
if (!allthreads || alive) out = arg.a * out;
9595
} else {
9696
SharedMemoryCache<Vector> cache {*this};
9797
Vector tmp;
98-
if (!allthreads || active) {
98+
if (!allthreads || alive) {
9999
// apply the chiral and flavor twists
100100
// use consistent load order across s to ensure better cache locality
101101
Vector x = arg.x[src_idx](my_flavor_idx, my_spinor_parity);
@@ -120,7 +120,7 @@ namespace quda
120120
// tmp += (c * tau_1) * x
121121
}
122122
cache.sync();
123-
if (!allthreads || active) {
123+
if (!allthreads || alive) {
124124
tmp += arg.c * cache.load_y(target::thread_idx().y + 1 - 2 * flavor);
125125

126126
// add the Wilson part with normalisation
@@ -132,7 +132,8 @@ namespace quda
132132
out = x + arg.a * out;
133133
}
134134

135-
if (active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
135+
if (!allthreads || alive)
136+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
136137
}
137138
};
138139

‎include/kernels/dslash_ndeg_twisted_clover_preconditioned.cuh‎

Lines changed: 56 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ namespace quda
1313
using WilsonArg<Float, nColor, nDim, DDArg, reconstruct_>::nSpin;
1414
static constexpr int length = (nSpin / (nSpin / 2)) * 2 * nColor * nColor * (nSpin / 2) * (nSpin / 2) / 2;
1515
static constexpr bool dynamic_clover = clover::dynamic_inverse();
16-
16+
1717
typedef typename mapper<Float>::type real;
1818
typedef typename clover_mapper<Float, length>::type C;
1919
const C A;
@@ -65,7 +65,7 @@ namespace quda
6565
out(x) = M*in = in + a*(C + i*b*gamma_5*tau_3 + c*tau_1)/(C^2 + b^2 - c^2)*D*x ( xpay == true )
6666
*/
6767
template <KernelType mykernel_type = kernel_type, bool allthreads = false>
68-
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool active = true)
68+
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool alive = true)
6969
{
7070
using namespace linalg; // for Cholesky
7171
typedef typename mapper<typename Arg::Float>::type real;
@@ -75,28 +75,30 @@ namespace quda
7575

7676
int src_idx = src_flavor / 2;
7777
int flavor = src_flavor % 2;
78+
bool active = mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
7879
int thread_dim; // which dimension is thread working on (fused kernel only)
7980
auto coord = getCoords<QUDA_4D_PC, mykernel_type>(arg, idx, flavor, parity, thread_dim);
8081

8182
const int my_spinor_parity = arg.nParity == 2 ? parity : 0;
8283
int my_flavor_idx = coord.x_cb + flavor * arg.dc.volume_4d_cb;
8384
Vector out;
84-
active &= mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
85-
if (arg.dd_out.isZero(coord)) {
86-
if (active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
87-
if constexpr (!allthreads) return;
88-
else active = false;
85+
if (!allthreads || alive) {
86+
if (arg.dd_out.isZero(coord)) {
87+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
88+
if constexpr (!allthreads) return;
89+
else alive = false;
90+
}
8991
}
9092

91-
if (!allthreads || active) {
93+
if (!allthreads || alive) {
9294
// defined in dslash_wilson.cuh
9395
applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
94-
}
9596

96-
if (mykernel_type != INTERIOR_KERNEL && active) {
97-
// if we're not the interior kernel, then we must sum the partial
98-
Vector x = arg.out[src_idx](my_flavor_idx, my_spinor_parity);
99-
out += x;
97+
if (mykernel_type != INTERIOR_KERNEL && active) {
98+
// if we're not the interior kernel, then we must sum the partial
99+
Vector x = arg.out[src_idx](my_flavor_idx, my_spinor_parity);
100+
out += x;
101+
}
100102
}
101103

102104
constexpr int n_flavor = 2;
@@ -109,7 +111,6 @@ namespace quda
109111

110112
int chirality = flavor; // relabel flavor as chirality
111113
SharedMemoryCache<HalfVector> cache {*this};
112-
113114
auto swizzle = [&](HalfVector x[2], int chirality) {
114115
if (chirality == 0)
115116
cache.save_y(x[1], target::thread_idx().y);
@@ -124,53 +125,57 @@ namespace quda
124125

125126
swizzle(out_chi, chirality); // apply the flavor-chirality swizzle between threads
126127

127-
if (isComplete<mykernel_type>(arg, coord) && active) {
128-
// load in the clover matrix
129-
HMat A = arg.A(coord.x_cb, parity, chirality);
128+
if (!allthreads || alive) {
129+
if (isComplete<mykernel_type>(arg, coord) && active) {
130+
// load in the clover matrix
131+
HMat A = arg.A(coord.x_cb, parity, chirality);
130132

131-
HalfVector A_chi[n_flavor];
133+
HalfVector A_chi[n_flavor];
132134
#pragma unroll
133-
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
134-
const complex<real> b(0.0, (chirality^flavor_) == 0 ? arg.b : -arg.b);
135-
A_chi[flavor_] = A * out_chi[flavor_];
136-
A_chi[flavor_] += b * out_chi[flavor_];
137-
A_chi[flavor_] += arg.c * out_chi[1 - flavor_];
138-
}
139-
140-
if constexpr (Arg::dynamic_clover) {
141-
HMat A2 = A.square();
142-
A2 += arg.b2_minus_c2;
143-
Cholesky<HMatrix, clover::cholesky_t<typename Arg::Float>, Arg::nColor * Arg::nSpin / 2> cholesky(A2);
135+
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
136+
const complex<real> b(0.0, (chirality^flavor_) == 0 ? arg.b : -arg.b);
137+
A_chi[flavor_] = A * out_chi[flavor_];
138+
A_chi[flavor_] += b * out_chi[flavor_];
139+
A_chi[flavor_] += arg.c * out_chi[1 - flavor_];
140+
}
141+
142+
if constexpr (Arg::dynamic_clover) {
143+
HMat A2 = A.square();
144+
A2 += arg.b2_minus_c2;
145+
Cholesky<HMatrix, clover::cholesky_t<typename Arg::Float>, Arg::nColor * Arg::nSpin / 2> cholesky(A2);
144146

145147
#pragma unroll
146-
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
147-
out_chi[flavor_] = static_cast<real>(0.25) * cholesky.backward(cholesky.forward(A_chi[flavor_]));
148-
}
149-
} else {
150-
HMat A2inv = arg.A2inv(coord.x_cb, parity, chirality);
148+
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
149+
out_chi[flavor_] = static_cast<real>(0.25) * cholesky.backward(cholesky.forward(A_chi[flavor_]));
150+
}
151+
} else {
152+
HMat A2inv = arg.A2inv(coord.x_cb, parity, chirality);
151153
#pragma unroll
152-
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
153-
out_chi[flavor_] = static_cast<real>(2.0) * (A2inv * A_chi[flavor_]);
154-
}
155-
}
154+
for (int flavor_ = 0; flavor_ < n_flavor; flavor_++) {
155+
out_chi[flavor_] = static_cast<real>(2.0) * (A2inv * A_chi[flavor_]);
156+
}
157+
}
158+
}
156159
}
157160

158161
swizzle(out_chi, chirality); // undo the flavor-chirality swizzle
159162

160-
if (isComplete<mykernel_type>(arg, coord) && active) {
161-
Vector tmp = out_chi[0].chiral_reconstruct(0) + out_chi[1].chiral_reconstruct(1);
162-
tmp.toNonRel(); // switch back to non-chiral basis
163-
164-
if (xpay && !arg.dd_x.isZero(coord)) {
165-
Vector x = arg.x[src_idx](my_flavor_idx, my_spinor_parity);
166-
out = x + arg.a * tmp;
167-
} else {
168-
// multiplication with a needed here?
169-
out = arg.a * tmp;
170-
}
163+
if (!allthreads || alive) {
164+
if (isComplete<mykernel_type>(arg, coord) && active) {
165+
Vector tmp = out_chi[0].chiral_reconstruct(0) + out_chi[1].chiral_reconstruct(1);
166+
tmp.toNonRel(); // switch back to non-chiral basis
167+
168+
if (xpay && !arg.dd_x.isZero(coord)) {
169+
Vector x = arg.x[src_idx](my_flavor_idx, my_spinor_parity);
170+
out = x + arg.a * tmp;
171+
} else {
172+
// multiplication with a needed here?
173+
out = arg.a * tmp;
174+
}
175+
}
176+
177+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
171178
}
172-
173-
if (active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
174179
}
175180
};
176181
} // namespace quda

‎include/kernels/dslash_ndeg_twisted_mass_preconditioned.cuh‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -65,37 +65,35 @@ namespace quda
6565
*/
6666

6767
template <KernelType mykernel_type = kernel_type, bool allthreads = false>
68-
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool active = true)
68+
__device__ __host__ __forceinline__ void operator()(int idx, int src_flavor, int parity, bool alive = true)
6969
{
7070
typedef typename mapper<typename Arg::Float>::type real;
7171
typedef ColorSpinor<real, Arg::nColor, 4> Vector;
7272

7373
int src_idx = src_flavor / 2;
7474
int flavor = src_flavor % 2;
7575

76-
bool active2 = mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
76+
bool active = mykernel_type != EXTERIOR_KERNEL_ALL; // is thread active (non-trival for fused kernel only)
7777
int thread_dim; // which dimension is thread working on (fused kernel only)
7878
auto coord = getCoords<QUDA_4D_PC, mykernel_type>(arg, idx, flavor, parity, thread_dim);
7979

8080
const int my_spinor_parity = arg.nParity == 2 ? parity : 0;
8181
int my_flavor_idx = coord.x_cb + flavor * arg.dc.volume_4d_cb;
8282
Vector out;
83-
if (!allthreads || active) {
83+
if (!allthreads || alive) {
8484
if (arg.dd_out.isZero(coord)) {
85-
if (mykernel_type != EXTERIOR_KERNEL_ALL) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
86-
if (!allthreads) return;
87-
active = false;
85+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
86+
if constexpr (!allthreads) return;
87+
else alive = false;
8888
}
8989
}
9090

91-
if (!allthreads || active) {
91+
if (!allthreads || alive) {
9292
if (!dagger || Arg::asymmetric) // defined in dslash_wilson.cuh
93-
applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active2, src_idx);
93+
applyWilson<dagger, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
9494
else // defined in dslash_twisted_mass_preconditioned
95-
applyWilsonTM<dagger, 2, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active2, src_idx);
96-
}
95+
applyWilsonTM<dagger, 2, mykernel_type>(out, arg, coord, parity, idx, thread_dim, active, src_idx);
9796

98-
if (!allthreads || active) {
9997
if (xpay && mykernel_type == INTERIOR_KERNEL && !arg.dd_x.isZero(coord)) {
10098
if constexpr (!dagger || Arg::asymmetric) { // apply inverse twist which is undone below
10199
// use consistent load order across s to ensure better cache locality
@@ -109,7 +107,7 @@ namespace quda
109107
Vector x = arg.x[src_idx](my_flavor_idx, my_spinor_parity);
110108
out += x; // just directly add since twist already applied in the dslash
111109
}
112-
} else if (mykernel_type != INTERIOR_KERNEL && active2) {
110+
} else if (mykernel_type != INTERIOR_KERNEL && active) {
113111
// if we're not the interior kernel, then we must sum the partial
114112
Vector x = arg.out[src_idx](my_flavor_idx, my_spinor_parity);
115113
out += x;
@@ -118,15 +116,15 @@ namespace quda
118116

119117
if constexpr (!dagger || Arg::asymmetric) { // apply A^{-1} to D*in
120118
SharedMemoryCache<Vector> cache {*this};
121-
if (!allthreads || active) {
122-
if (isComplete<mykernel_type>(arg, coord) && active2) {
119+
if (!allthreads || alive) {
120+
if (isComplete<mykernel_type>(arg, coord) && active) {
123121
// to apply the preconditioner we need to put "out" in shared memory so the other flavor can access it
124122
cache.save(out);
125123
}
126124
}
127125
cache.sync(); // safe to sync here since other threads will exit if allowed, or all be here
128-
if (!allthreads || active) {
129-
if (isComplete<mykernel_type>(arg, coord) && active2) {
126+
if (!allthreads || alive) {
127+
if (isComplete<mykernel_type>(arg, coord) && active) {
130128
if (flavor == 0)
131129
out = arg.a * (out + arg.b * out.igamma(4) + arg.c * cache.load_y(target::thread_idx().y + 1));
132130
else
@@ -135,7 +133,8 @@ namespace quda
135133
}
136134
}
137135

138-
if (active && active2) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
136+
if (!allthreads || alive)
137+
if (mykernel_type != EXTERIOR_KERNEL_ALL || active) arg.out[src_idx](my_flavor_idx, my_spinor_parity) = out;
139138
}
140139
};
141140

0 commit comments

Comments
 (0)