Torch.where is too slow

AdilZouitine: <p>Hello everyone,<br>
I am doing reinforcement learning using a policy gradient algo.<br>
However I have impossible actions, therefore I modify the logits of my impossible actions before using categorical.</p>
<p>When applying this mask I use a torch.where.<br>
I have profiled the code and one of the biggest bottlenecks comes from this.<br>
Do you have any optimization to do to make it faster?</p>
<p>Best regards</p>
<pre><code class="lang-python">
import torch
from torch.distributions import Categorical


class CategoricalMasked(Categorical):
    LOGIT_INF_VALUE = -1E10

    def __init__(self, mask, logits):
        # dtype of mask torch.float32
        # size example logit and mask (4096, 5)
        self.mask = mask
        self.mask = torch.gt(self.mask,0.5)

        logits = torch.where(self.mask, logits,
                             torch.ones_like(logits) * torch.tensor(CategoricalMasked.LOGIT_INF_VALUE, device=logits.device)
        super(CategoricalMasked, self).__init__(None, logits, None)

    def entropy(self):
        p_log_p = self.logits * self.probs
        p_log_p = torch.where(self.mask, p_log_p, torch.tensor(0., device=self.logits.device)
        return -p_log_p.sum(-1)

</code></pre>

albanD: <p>Hi,</p>
<p>If you run on GPU, the following might be faster:</p>
<pre><code class="lang-python">self.mask = self.mask.to(logits.dtype)
logits = self.mask * logits + (1 - self.mask) * CategoricalMasked.LOGIT_INF_VALUE
</code></pre>

AdilZouitine: <p>Hi <a class="mention" href="/u/alband">@albanD</a> good news!<br>
Why  <code>torch.where</code> is slower than what you suggest ?</p>

googlebot: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="100915">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<pre><code class="lang-auto">logits = self.mask * logits + (1 - self.mask) * CategoricalMasked.LOGIT_INF_VALUE
</code></pre>
</blockquote>
</aside>
<p>that shouldn’t be faster than: torch.lerp(LOGIT_INF_TENSOR, logits, self.mask)</p>
<aside class="quote no-group" data-username="AdilZouitine" data-post="1" data-topic="100915">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/adilzouitine/48/8345_2.png" class="avatar"> AdilZouitine:</div>
<blockquote>
<pre><code class="lang-auto">p_log_p = torch.where(self.mask, p_log_p, torch.tensor(0., device=self.logits.device)
</code></pre>
</blockquote>
</aside>
<p>p_log_p *= self.mask</p>

AdilZouitine: <p>Thanks II’l try it ! Just one question why where is slower than lerp ? <a class="mention" href="/u/googlebot">@googlebot</a></p>

googlebot: <p>lerp should just fuse <code>x1*w+x2*(1-w)</code> expressions, like the one <a class="mention" href="/u/alband">@albanD</a> suggested, avoiding memory allocations</p>
<p>I don’t know whether this or torch.where is faster. The latter could be slower if it screws memory access patterns on cuda (but this is avoidable, I’d think)</p>

albanD: <p>Ho I don’t have the answer for which one is faster. But I know that point-wise addition/multiplication kernels are heavily optimized and being worked on. While the <code>where</code> kernel is not scrutinized as much. But curious to know which one is faster actually compared to Alex’s proposal as well.</p>