Compiling stack of Conv3d increases runtime with redundant format conversions

ajayj: <p>On both PyTorch 2.3.1+cu121 (Colab A100) and 2.4.0+cu124 (H100), torch.compile is introducing unnecessary memory layout change kernels around Conv3d layers. These Conv3d layers have 3x3x3 kernel size, 1x1x1 padding, no bias, no activation and preserve the channel dimension.</p>
<pre><code class="lang-auto">class ConvNet(nn.Module):
    """Simple convolutional network."""
    def __init__(self, in_channels: int, hidden_channels: int, out_channels: int):
        super().__init__()
        self.conv_1 = nn.Conv3d(in_channels, hidden_channels, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        self.conv_2 = nn.Conv3d(hidden_channels, hidden_channels, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        self.conv_3 = nn.Conv3d(hidden_channels, out_channels, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)

    def forward(self, x) -&gt; Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass.
        Args:
            x: Input 3D tensor. Shape: [B, C, T, H, W].
        """
        assert x.ndim == 5

        with record_function("conv_layers"):
            x = self.conv_1(x)
            x = self.conv_2(x)
            x = self.conv_3(x)

        assert x.is_contiguous(memory_format=torch.channels_last_3d)

        return x
</code></pre>
<p>Without compilation, the stack takes 6.482 ms / step after warmup on an input with shape (4, 8, 64, 256, 256) and 64 channels in the network.</p>
<p>When torch.compiled, the same network takes 10.878 ms / step.</p>
<p>The PyTorch compiler is introducing an unnecessary nhwcToNchwKernel immediately before an nchwToNhwcKernel between every pair of Conv3d layers:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61774844912c33671b43917553e0699c749feb94.png" data-download-href="https://discuss.pytorch.org/uploads/default/61774844912c33671b43917553e0699c749feb94" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61774844912c33671b43917553e0699c749feb94_2_690x141.png" alt="image" data-base62-sha1="dUdSitCkUfi3bgHQgW8YBpm9i4s" width="690" height="141" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61774844912c33671b43917553e0699c749feb94_2_690x141.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61774844912c33671b43917553e0699c749feb94_2_1035x211.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61774844912c33671b43917553e0699c749feb94_2_1380x282.png 2x" data-dominant-color="DCE1E5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2042×418 41.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>This leads to a 1.68x slowdown on A100.</p>
<p>Minimal reproducible example. The input tensor and the network are in channels_last_3d format. <a href="https://colab.research.google.com/drive/18TT95jUtg6XrBP6TkwSZC8BsiI8amtp7?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a></p>
<p>Profiles and code are here: <a href="https://drive.google.com/drive/folders/1nw78W4dPPh9uQxxOvzWkeaGk8LIfi6pB?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Conv3d bug - Google Drive</a></p>

ptrblck: <p>CC <a class="mention" href="/u/marksaroufim">@marksaroufim</a> and <a class="mention" href="/u/eqy">@eqy</a> - would you know if <code>torch.compile</code> tries to use the channels-last format by default for these layers and adds these transformations?</p>

marksaroufim: <p><a class="mention" href="/u/eellison">@eellison</a> is the lead for inductor and would likely know</p>

Elias_Ellison: <p>Yes, we do attempt to use channels last by default depending on the convs seen. You can disable with <code>TORCHINDUCTOR_LAYOUT_OPTIMIZATION=0</code> or <code>torch._inductor.config.layout_optimization = False</code>.</p>
<p>Maybe we need to be less agressive with conv3d. Would you mind filing issue ?</p>
<p>Separately - is there a reason you are using autocast instead of converting parameters to bfloat16 before running model ?</p>

eqy: <p>I believe I’ve been forwarded this issue previously and have passed it on to cuDNN. To comment I’m not sure <code>compile</code> is really to blame as even in eager mode we observed the format conversions</p>