How to add input shape dimension hint to an optimized graph?

username: masahi
cooked: <p>Hi, I’m trying to get an optimized graph from graph executor via ScriptModule’s <code>graph_for(...)</code> method. A simple test case is below:</p>
<pre><code class="lang-auto">import torch

conv = torch.nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3)

#  To avoid dealing with prim::Bailout stuff
torch._C._jit_set_profiling_executor(False)

inp = torch.rand(1, 3, 224, 224)
trace = torch.jit.trace(conv, inp).eval()

print(trace.graph_for(inp))
</code></pre>
<p>But it seems graph executor erases shape information on construction, so the output of <code>graph_for</code> has shape information removed.</p>
<pre><code class="lang-auto">graph(%self : __torch__.torch.nn.modules.module.Module,
      %input : Float(*, *, *, *)):
  %23 : int[] = prim::Constant[value=[0, 0]]()
  %22 : int[] = prim::Constant[value=[1, 1]]()
  %4 : int = prim::Constant[value=1]() # /home/masa/projects/deep/pytorch/torch/nn/modules/conv.py:345:0
  %13 : bool = prim::Constant[value=0]() # /home/masa/projects/deep/pytorch/torch/nn/modules/conv.py:345:0
  %20 : bool = prim::Constant[value=1]() # /home/masa/projects/deep/pytorch/torch/nn/modules/conv.py:345:0
  %2 : Float(*) = prim::GetAttr[name="bias"](%self)
  %3 : Float(*, *, *, *) = prim::GetAttr[name="weight"](%self)
  %21 : Float(*, *, *, *) = aten::_convolution(%input, %3, %2, %22, %23, %22, %13, %23, %4, %13, %13, %20) # /home/masa/projects/deep/pytorch/torch/nn/modules/conv.py:345:0
  return (%21)

</code></pre>
<p>If I know that my input shape is always fixed, is there a way to add explicit shape information? I know that the output of trace has shape information preserved, but I want an optimized graph available via <code>graph_for(...)</code>.</p>

username: villedepommes
cooked: <p>Could you please elaborate on your use case? The exact static shapes were added to enable backends to generate more efficient kernels. I’m not sure they will be added to serialization or any other public interface.</p>
<p>WIthout,<br>
<code>torch._C._jit_set_profiling_executor(False)</code> we don’t capture the exact shapes only tensor ranks. Even the rank information, I believe, considered internal for use by backends and optimization passes.</p>

username: masahi
cooked: <p>ok, my use case is to translate Torchscript IR to <a href="https://github.com/apache/incubator-tvm" rel="nofollow noopener">TVM</a> compiler’s IR. I have a repo <a href="https://github.com/masahi/torchscript-to-tvm" rel="nofollow noopener">https://github.com/masahi/torchscript-to-tvm</a> which demonstrates translating torch models to TVM, compile and run it under TVM and get the identical output as torch.</p>
<p>Right now I take the output of <code>torch.jit.trace(...)</code> which has shape information attached, and for each torch operator node I translate it to corresponding one in TVM. Since TVM is mostly a static compiler, shape information is required.</p>
<p>TVM has its own set of optimization passes, so it is no problem to take the unoptimized input torch IR. Currently I’m applying the Torchscript inlining pass to remove <code>prim::CallMethod</code> wrapping,  but it seems rather ad hoc to me and I rather want to apply other optimization passes in torch as well.</p>
<p>I know I can apply each optimization passes manually (and I do for the inline pass), but the API prefix <code>torch._C._jit_pass*</code> suggests they are not “officially” supported, so I’m not sure if I want to use them directly. Since I discovered that I can access the optimized graph via <code>graph_for(...)</code> method, I’m looking to see if this is something I can use.</p>
<p>I disabled the profiling executor because it adds <code>prim::Bailout</code> and <code>prim::BailoutTemplate</code> nodes for which I have no idea how to translate to TVM. Since input shape is static in TVM, I think they are not relevant to my use case, so I don’t want to see them in the input IR.</p>

username: eellison
cooked: <p>I would suggest that you run it with the profiling executor a few times with inputs that cover the different Tensor dimensions you expect to use, and then add a pass to remove Bailout nodes. This should give you a Graph with shape information to maximum generality.</p>

username: eellison
cooked: <p>There are also a few passes that will be landed shortly in Pytorch JIT that should help with this conversion. Freezing and a functionalization pass. I will comment here when they are landed.</p>

username: masahi
cooked: <p>Thanks, if removing bailout is possible that would definitely work for me. At a quick glance there is no pass to remove Bailout nodes in torch, but I’ll try if I can do it from “userland”.</p>
<p>Another API I’m interested in is <code>_propagate_shapes</code> function:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/8527ba8b7059d8c7df6fc2c11271ceefbf91ce5a/torch/csrc/jit/script/init.cpp#L476-L485" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/8527ba8b7059d8c7df6fc2c11271ceefbf91ce5a/torch/csrc/jit/script/init.cpp#L476-L485" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/8527ba8b7059d8c7df6fc2c11271ceefbf91ce5a/torch/csrc/jit/script/init.cpp#L476-L485</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="476" style="counter-reset: li-counter 475 ;">
<li>static std::shared_ptr&lt;Graph&gt; _propagate_shapes(</li>
<li>    Graph&amp; graph,</li>
<li>    std::vector&lt;at::Tensor&gt; inputs,</li>
<li>    bool with_grad = false) {</li>
<li>  Stack stack(inputs.begin(), inputs.end());</li>
<li>  auto retval = graph.copy();</li>
<li>  setInputTensorTypes(*retval, stack, /*complete=*/false);</li>
<li>  PropagateInputShapes(retval);</li>
<li>  return retval;</li>
<li>}</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Do you know the use case of this function and if this is something I should take a look?</p>

username: eellison
cooked: <p>This API is only useful with the executor that isn’t the profiled executor, so I don’t think it applies. Yes, the pass doesn’t exactly exist as you need it. The logic should pretty much be the same as <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/guard_elimination.cpp#L145" rel="nofollow noopener">right here</a>, except you are always removing the guards.</p>