Torch.compile cpp backend not fusing OuterLoopFusedSchedulerNode properly

vivekvpandya: <p>I have following simple test case for which I see OuterLoopFusedSchedulerNode generated</p>
<pre><code class="lang-auto">def test_group_norm():
    from torch._inductor import config

    # Force multple scheduler nodes creation to fuse them
    config.realize_opcount_threshold = 1
    config.inplace_buffers = False
    class M(torch.nn.Module):
        def __init__(self) -&gt; None:
            super().__init__()
            self.group_norm = torch.nn.GroupNorm(2, 16)

        def forward(self, x):
            return self.group_norm(x)

    mod = M().eval()
    x = torch.randn(2, 16, 32, 32)
    with torch.no_grad():
        expected = mod(x)
        compiled_m = torch.compile(mod, backend="inductor", fullgraph=True)
        actual = compiled_m(x)
        assert torch.allclose(
            expected.cpu(), actual.cpu(), rtol=0.000001, atol=0.001
        ), "Test failed: Output does not match expected output"
</code></pre>
<p>However in generated code we can see that loops like following should have been fused to single outer loop</p>
<pre><code class="lang-auto">for(long x0=static_cast&lt;long&gt;(0L); x0&lt;static_cast&lt;long&gt;(4L); x0+=static_cast&lt;long&gt;(1L))
</code></pre>
<p>so why this does not happen? is this intentional or a bug?</p>
<pre data-code-wrap="cpp"><code class="lang-cpp">extern "C" void kernel(const float* in_ptr0,
                       const float* in_ptr1,
                       const float* in_ptr2,
                       float* out_ptr0,
                       float* out_ptr1,
                       float* out_ptr4,
                       float* out_ptr6,
                       float* out_ptr8)
{
    // This is outerloop fused node buf0_buf1_buf3_buf5_buf6_buf4_buf7
    {
        #pragma GCC ivdep
        for(long x0=static_cast&lt;long&gt;(0L); x0&lt;static_cast&lt;long&gt;(4L); x0+=static_cast&lt;long&gt;(1L))
        {
            {
                Welford&lt;float&gt; tmp_acc0 = Welford&lt;float&gt;();
                Welford&lt;at::vec::Vectorized&lt;float&gt;&gt; tmp_acc0_vec = Welford&lt;at::vec::Vectorized&lt;float&gt;&gt;();
                static WeightRecp&lt;at::vec::Vectorized&lt;float&gt;&gt; weight_recps(static_cast&lt;long&gt;(512L));
                for(long x1=static_cast&lt;long&gt;(0L); x1&lt;static_cast&lt;long&gt;(8192L); x1+=static_cast&lt;long&gt;(16L))
                {
                    auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(in_ptr0 + static_cast&lt;long&gt;(x1 + (8192L*x0)), 16);
                    tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp0, &amp;weight_recps);
                }
                tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(tmp_acc0_vec));
                out_ptr0[static_cast&lt;long&gt;(x0)] = static_cast&lt;float&gt;(tmp_acc0.mean);
                out_ptr1[static_cast&lt;long&gt;(x0)] = static_cast&lt;float&gt;(tmp_acc0.m2);
            }
        }
    }
    {
        #pragma omp simd simdlen(8) 
        for(long x0=static_cast&lt;long&gt;(0L); x0&lt;static_cast&lt;long&gt;(4L); x0+=static_cast&lt;long&gt;(1L))
        {
            auto tmp0 = out_ptr1[static_cast&lt;long&gt;(x0)];
            auto tmp1 = static_cast&lt;float&gt;(8192.0);
            auto tmp2 = tmp0 / tmp1;
            auto tmp3 = static_cast&lt;float&gt;(1e-05);
            auto tmp4 = decltype(tmp2)(tmp2 + tmp3);
            auto tmp5 = 1 / std::sqrt(tmp4);
            out_ptr4[static_cast&lt;long&gt;(x0)] = tmp5;
        }
    }
    {
        #pragma GCC ivdep
        for(long x0=static_cast&lt;long&gt;(0L); x0&lt;static_cast&lt;long&gt;(4L); x0+=static_cast&lt;long&gt;(1L))
        {
            for(long x1=static_cast&lt;long&gt;(0L); x1&lt;static_cast&lt;long&gt;(8192L); x1+=static_cast&lt;long&gt;(16L))
            {
                auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(in_ptr0 + static_cast&lt;long&gt;(x1 + (8192L*x0)), 16);
                auto tmp1 = out_ptr0[static_cast&lt;long&gt;(x0)];
                auto tmp4 = out_ptr4[static_cast&lt;long&gt;(x0)];
                auto tmp2 = at::vec::Vectorized&lt;float&gt;(tmp1);
                auto tmp3 = tmp0 - tmp2;
                auto tmp5 = at::vec::Vectorized&lt;float&gt;(tmp4);
                auto tmp6 = tmp3 * tmp5;
                tmp6.store(out_ptr6 + static_cast&lt;long&gt;(x1 + (8192L*x0)));
            }
        }
    }
    // This is Fused node buf8_buf9
    {
        #pragma GCC ivdep
        for(long x0=static_cast&lt;long&gt;(0L); x0&lt;static_cast&lt;long&gt;(2L); x0+=static_cast&lt;long&gt;(1L))
        {
            #pragma GCC ivdep
            for(long x1=static_cast&lt;long&gt;(0L); x1&lt;static_cast&lt;long&gt;(16L); x1+=static_cast&lt;long&gt;(1L))
            {
                for(long x2=static_cast&lt;long&gt;(0L); x2&lt;static_cast&lt;long&gt;(1024L); x2+=static_cast&lt;long&gt;(16L))
                {
                    auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(out_ptr6 + static_cast&lt;long&gt;(x2 + (1024L*x1) + (16384L*x0)), 16);
                    auto tmp1 = in_ptr1[static_cast&lt;long&gt;(x1)];
                    auto tmp4 = in_ptr2[static_cast&lt;long&gt;(x1)];
                    auto tmp2 = at::vec::Vectorized&lt;float&gt;(tmp1);
                    auto tmp3 = tmp0 * tmp2;
                    auto tmp5 = at::vec::Vectorized&lt;float&gt;(tmp4);
                    auto tmp6 = tmp3 + tmp5;
                    tmp6.store(out_ptr8 + static_cast&lt;long&gt;(x2 + (1024L*x1) + (16384L*x0)));
                }
            }
        }
    }
}
</code></pre>

leslie-fang-intel: <p>Hi <a class="mention" href="/u/vivekvpandya">@vivekvpandya</a>, thanks for starting the discussion!</p>
<p>Before diving into my understanding of the issue, I’d like to briefly explain how we generate code for <code>OuterLoopFusedSchedulerNode</code> in the Inductor C++ backend:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/159e2f96e3d7aec0187b03c483e177c5e96b1ebb/torch/_inductor/codegen/cpp.py#L5045">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/159e2f96e3d7aec0187b03c483e177c5e96b1ebb/torch/_inductor/codegen/cpp.py#L5045" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/159e2f96e3d7aec0187b03c483e177c5e96b1ebb/torch/_inductor/codegen/cpp.py#L5045" target="_blank" rel="noopener nofollow ugc">torch/_inductor/codegen/cpp.py</a></h4>

<div class="git-blob-info">
  <a href="https://github.com/pytorch/pytorch/blob/159e2f96e3d7aec0187b03c483e177c5e96b1ebb/torch/_inductor/codegen/cpp.py#L5045" rel="noopener nofollow ugc"><code>159e2f96e</code></a>
</div>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="5035" style="counter-reset: li-counter 5034 ;">
          <li>def codegen_node(</li>
          <li>    self,</li>
          <li>    node: Union[OuterLoopFusedSchedulerNode, FusedSchedulerNode, SchedulerNode],</li>
          <li>):</li>
          <li>    """</li>
          <li>    Turn an set of pre-fused nodes into a C++ kernel.</li>
          <li>    """</li>
          <li>    kernel_group = self.kernel_group</li>
          <li></li>
          <li>    if isinstance(node, OuterLoopFusedSchedulerNode):</li>
          <li class="selected">        self.codegen_outer_loop_node(node)</li>
          <li>    else:</li>
          <li>        nodes: list[SchedulerNode] = node.get_nodes()  # type: ignore[assignment]</li>
          <li>        nodes = self.try_loop_split(nodes)</li>
          <li>        cpp_kernel_proxy = CppKernelProxy(kernel_group)</li>
          <li>        cpp_kernel_proxy.codegen_nodes(nodes)</li>
          <li>        kernel_group.finalize_kernel(cpp_kernel_proxy, nodes)</li>
          <li></li>
          <li>    args_num = self._get_scheduled_num_args()</li>
          <li>    if args_num &gt; CppScheduling.MAX_FUSED_KERNEL_ARGS_NUM:</li>
          <li>        self._set_flush_status(True)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>We first generate code for each individual <code>SchedulerNode</code> within the <code>OuterLoopFusedSchedulerNode</code>. After that, we check whether any loop levels from the generated kernels can be merged at the outer loop level. Specifically, we examine whether two loop levels from different kernels share the same loop attributes—such as whether the loop is vectorized and the vectorization step size. If the attributes match, we merge the loop levels.</p>
<p>Looking at the generated code you posted, it seems there are three loop levels. The middle one:</p>
<pre><code class="lang-auto">#pragma omp simd simdlen(8) 
for (long x0 = static_cast&lt;long&gt;(0L); x0 &lt; static_cast&lt;long&gt;(4L); x0 += static_cast&lt;long&gt;(1L))
</code></pre>
<p>appears to have different loop attributes compared to the loop level before and after it. That might be the reason why outer loop fusion didn’t take effect.</p>
<p>If needed, I’m happy to take a closer look—please share the PyTorch version, CPU info, and the command you used to run it.</p>

vivekvpandya: <p>Thanks <a class="mention" href="/u/leslie-fang-intel">@leslie-fang-intel</a> for quicky reply. I am on pytorch 2.4 with CPU Model name:             AMD EPYC 9124 16-Core Processor.</p>
<p>I debugged and found that node is marked as OuterLoopFusedSchedulerNode but it fails to merge outer loop and fall backs to normal code gen at <a href="https://github.com/pytorch/pytorch/blob/ee1b6804381c57161c477caa380a840a84167676/torch/_inductor/codegen/cpp.py#L3787" class="inline-onebox" rel="noopener nofollow ugc">pytorch/torch/_inductor/codegen/cpp.py at ee1b6804381c57161c477caa380a840a84167676 · pytorch/pytorch · GitHub</a></p>
<p>I also checked with torch 2.6 it also fails to merge outer loop.</p>