`masked_scatter` slower than `where`?

username: dengcy
cooked: <p>If I want to replace a tensor <code>x</code> partially by <code>y</code> based on a <code>condition</code> (a Boolean tensor), I think the most suitable method is to call <code>masked_scatter</code> since it is designed for such usage. However, I found <code>torch.where</code> is even faster?</p>
<pre><code class="lang-auto">import torch
import cProfile

n=100000
dim=2000
device='cuda'

# masked_scatter_
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n, dim, device=device) &gt; 0.5
cProfile.run('x.masked_scatter_(condition, y)')


# torch.where
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n, dim, device=device) &gt; 0.5
cProfile.run('x = torch.where(condition, x, y)')
</code></pre>
<p>What I got is</p>
<pre><code class="lang-auto"> 4 function calls in 0.112 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.112    0.112 &lt;string&gt;:1(&lt;module&gt;)
        1    0.000    0.000    0.112    0.112 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
        1    0.112    0.112    0.112    0.112 {method 'masked_scatter_' of 'torch._C._TensorBase' objects}


4 function calls in 0.000 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.000    0.000 &lt;string&gt;:1(&lt;module&gt;)
        1    0.000    0.000    0.000    0.000 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {built-in method where}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
</code></pre>
<p>Using CPU has a similar result. Besides, <code>torch.where</code> has a lower memory usage (not shown in the result above) despite I am using a inplace <code>masked_scatter_</code>. Why would this happen?</p>

username: ptrblck
cooked: <p>Note that CUDA operations are executed asynchronously, so you would have to synchronize the code before starting and stopping timers.<br>
I’m not deeply familiar with <code>cProfile</code>, but assume that it doesn’t synchronize the GPU internally.<br>
The <code>torch.utils.benchmark</code> utils might be useful to compare different methods, as they will add warmup iterations and synchronize internally.</p>

username: dengcy
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="115910">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>torch.utils.benchmark</p>
</blockquote>
</aside>
<p>Thanks for your reply. I tried <code>torch.utils.benchmark</code> but I got similar results</p>
<pre><code class="lang-auto">import torch
n=100000
dim=2000
device='cuda'
import torch.utils.benchmark as benchmark

# masked_scatter_
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n,1, device=device) &gt; 0.5
print(benchmark.Timer('x.masked_scatter_(condition.expand(n,dim), y)',globals={'x':x,'y':y,'condition':condition,'n':n,'dim':dim}).timeit(10))

# torch.where
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n,1, device=device) &gt; 0.5
print(benchmark.Timer('torch.where(condition.expand(n,dim), x, y)',globals={'x':x,'y':y,'condition':condition,'n':n,'dim':dim}).timeit(10))
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">&lt;torch.utils.benchmark.utils.common.Measurement object at 0x0000019B0E340670&gt;
x.masked_scatter_(condition.expand(n,dim), y)
  22.45 ms
  1 measurement, 10 runs , 1 thread
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x0000019B097ADCA0&gt;
torch.where(condition.expand(n,dim), x, y)
  5.53 ms
  1 measurement, 10 runs , 1 thread
</code></pre>

username: theevann
cooked: <p>Hello,</p>
<p>I have no answer but I made a similar observation today.<br>
I compared tensor indexing vs <code>masked_scatter</code>:</p>
<pre><code class="lang-python">x[condition] = y[condition]
# VS
x.masked_scatter_(condition.view(-1, 1), y)
</code></pre>
<p>and indeed <code>masked_scatter</code> is much slower.</p>
<p>Additionally, I found the solution using tensor indexing to be faster than <code>torch.where</code> by about 20%, and to be more readable.</p>
<hr>
<p>On my machine, I get:</p>
<pre><code class="lang-auto">&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f7845fc4b70&gt;
x.masked_scatter_(condition.view(-1,1), y)
  28.40 ms
  1 measurement, 100 runs , 1 thread

&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f7845fc4ac8&gt;
torch.where(condition.view(-1,1), x, y)
  6.19 ms
  1 measurement, 100 runs , 1 thread

&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f7845fc4828&gt;
x[condition] = y[condition]
  4.79 ms
  1 measurement, 100 runs , 1 thread
</code></pre>

username: dengcy
cooked: <p>Yes, basically the same as what I got. Comparing <code>where</code> and <code>x[condition] = y[condition]</code>, the advantage of <code>where</code> is that it has lower memory usage. I do not know how to get the benchmark, but I observed from my task manager.</p>