Understanding Model Compilation/Optimization

johnc1231: <p>I’m experimenting with the idea of generating Torchscript IR as part of a project I’m working on, and would like to better understand how it’s compiled and optimized.</p>
<p>I’ve followed the tutorial here: <a href="https://pytorch.org/tutorials/advanced/cpp_export.html" class="inline-onebox" rel="noopener nofollow ugc">Loading a TorchScript Model in C++ — PyTorch Tutorials 1.7.1 documentation</a> and successfully loaded a Torchscript model into C++ and ran it on some example data. I’ve also found the overview here: <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/OVERVIEW.md#python-compiler-interaction" class="inline-onebox" rel="noopener nofollow ugc">pytorch/OVERVIEW.md at master · pytorch/pytorch · GitHub</a> helpful.</p>
<p>What I’d like to know is:</p>
<ol>
<li>
<p>When I load a module in in C++, is there any compilation / optimization that can performed before the first time I call the <code>forward</code> method? Or do optimizations only happen with a JIT that runs after data is first passed in?</p>
</li>
<li>
<p>If there are JIT optimizations to be done when I call the <code>forward</code> method, do the optimizations persist so that next time I call the <code>forward</code> method on the same module object, they won’t have to be performed again?</p>
</li>
<li>
<p>If a module/model gets optimized by the JIT, do any of those optimizations change the output I’d get if I now wrote the model back out to a torchscript <code>.pt</code> file?</p>
</li>
<li>
<p>Where does kernel fusion happen? Is that part of the JIT, or is that something that has to be figured out when generating the IR in python?</p>
</li>
<li>
<p>Does Pytorch have the ability to dynamically decide whether an operation should be run on CPU or GPU based on things like data size?</p>
</li>
</ol>
<p>If there’s some resource I haven’t yet seen that could help me better understand this stuff, a pointer towards it would be appreciated. Thanks for your help.</p>

johnc1231: <p>I’m happy to break these questions up into different posts if that would be helpful</p>

eellison: <p>Hi johnc1231,</p>
<ol>
<li>
<p>Not unless you invoke it. Currently the only AOT optimization is <a href="https://pytorch.org/docs/master/generated/torch.jit.freeze.html#torch-jit-freeze" class="inline-onebox" rel="noopener nofollow ugc">torch.jit.freeze — PyTorch master documentation</a>. It does a limited set of optimizations on 1.7, <a href="https://github.com/pytorch/pytorch/pull/50074" rel="noopener nofollow ugc">a little more on master</a>, and I’m working on adding more AOT stuff. I will add a C++ api equivalent by 1.8.</p>
</li>
<li>
<p>Yes, they do persist, however if you pass in tensors of a different type we may respecialize.</p>
</li>
<li>
<p>JIT optimizations will not affect saving your model.</p>
</li>
<li>
<p>That happens as part of the JIT. It’s happens after we profile the tensors that are seen at runtime.</p>
</li>
<li>
<p>No, that one of the main design points of pytorch is that it does not automatically decide what parts of the module to run on gpu vs cpu, for both eager &amp; jit. That is controlled by the user.</p>
</li>
</ol>
<p>To better understand optimizations I would suggest running a simple file like:</p>
<pre><code class="lang-auto">import torch

@torch.jit.script
def foo(x):
    return x + x + x

foo(torch.rand([4, 4], device='cuda'))
foo(torch.rand([4, 4], device='cuda'))
</code></pre>
<p>with <code>PYTORCH_JIT_LOG_LEVEL='profiling_graph_executor_impl' foo.py</code></p>

johnc1231: <p>Thanks for a very helpful response, and for planning to add freeze to the C++ API. I’ll experiment with the JIT log level printouts.</p>

johnc1231: <p><a class="mention" href="/u/eellison">@eellison</a> I did what you suggested, and I noticed the final optimized output looks like:</p>
<pre><code class="lang-auto">[DUMP profiling_graph_executor_impl.cpp:620] Optimized Graph:
[DUMP profiling_graph_executor_impl.cpp:620] graph(%x.1 : Tensor):
[DUMP profiling_graph_executor_impl.cpp:620]   %1 : int = prim::Constant[value=1]()
[DUMP profiling_graph_executor_impl.cpp:620]   %4 : Tensor = aten::add(%x.1, %x.1, %1) # triple_add.py:5:11
[DUMP profiling_graph_executor_impl.cpp:620]   %7 : Tensor = aten::add(%4, %x.1, %1) # triple_add.py:5:11
[DUMP profiling_graph_executor_impl.cpp:620]   return (%7)
</code></pre>
<p>It still makes two separate calls to <code>add</code>. I would think part of kernel fusion is that it could somehow make one call to <code>add</code> that takes in all three inputs. Am I misinterpreting the graph, or am I misunderstanding what kernel fusion does?</p>

eellison: <p>Were you using CUDA or CPU ? Could you post a more complete repro? CPU fusion is in progress &amp; targeting 1.9</p>

johnc1231: <p>Ah yeah, I was using CPU, that would explain it. I didn’t realize that GPU fusion was working but CPU fusion was in progress. Thank you.</p>