Skip to content

Commit 0c1227b

Browse files
author
Michael Abbott
committed
padding conditional store, without else clause
As suggested in JuliaSIMD/LoopVectorization.jl#174
1 parent 87358aa commit 0c1227b

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

src/macro.jl

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -466,7 +466,7 @@ padmodclamp_replace(s, store, inside=false) = s
466466
padmodclamp_replace(ex::Expr, store, inside=false) =
467467
if ex.head == :(=) && @capture_(ex.args[1], A_[inds__])
468468
# This tricky case is 𝛥A[pad(i,2)] = 𝛥A[pad(i,2)] + ...
469-
Aex, fun = padmodclamp_pair(A, inds, store)
469+
Aex, fun = padmodclamp_pair(A, inds, store, true)
470470
right = if fun != identity
471471
padmodclamp_replace(ex.args[2], store, true)
472472
else
@@ -481,7 +481,7 @@ padmodclamp_replace(ex::Expr, store, inside=false) =
481481
Expr(ex.head, args...)
482482
end
483483

484-
padmodclamp_pair(A, inds, store) = begin
484+
padmodclamp_pair(A, inds, store, assign=false) = begin
485485
nopadif = []
486486
inds4 = map(enumerate(inds)) do (d,ex)
487487
isexpr(ex, :call) || return ex
@@ -509,7 +509,9 @@ padmodclamp_pair(A, inds, store) = begin
509509
for c2 in nopadif[2:end]
510510
cond = :($cond & $c2)
511511
end
512-
if store.padkeyword == TYP # default
512+
if assign # for gradients, this wraps 𝛥A[pad(i,2)] = 𝛥A[pad(i,2)] + ...
513+
ex -> :($cond && $ex)
514+
elseif store.padkeyword == TYP # default, pad with zero
513515
ex -> :($cond ? $ex : zero(eltype($A)))
514516
else
515517
ex -> :($cond ? $ex : $convert($eltype($A), $(store.padkeyword)))

0 commit comments

Comments
 (0)