How to gracefully mask CompositeImplicitAutograd for different backends

yangxiaorun: <h1><a name="how-to-gracefully-mask-compositeimplicitautograd-for-different-backends-1" class="anchor" href="#how-to-gracefully-mask-compositeimplicitautograd-for-different-backends-1"></a>How to gracefully mask CompositeImplicitAutograd for different backends</h1>
<h2><a name="background-2" class="anchor" href="#background-2"></a>Background:</h2>
<p>I implemented torch.compile’s backend for my hardware via privateUserOne. I also found that torch.compile by default decomposes <code>upsample_nearest2d</code> into a bunch of small operators, just like <code>_upsample_nearest</code> does. But on my hardware, the <code>_unsafe_index</code> operator doesn’t perform well, so I’d like to be able to call the custom <code>upsample_nearest2d</code> operator directly for better performance.</p>
<pre><code class="lang-auto">@register_decomposition([aten.upsample_nearest2d.default, aten.upsample_nearest2d.out])
@aten.upsample_nearest2d.default.py_impl(DispatchKey.CompositeImplicitAutograd)
@aten.upsample_nearest2d.default.py_impl(DispatchKey.Autograd)
@out_wrapper(preserve_memory_format=True, exact_dtype=True)
def upsample_nearest2d(
    input: Tensor,
    output_size: List[int],
    scales_h: Optional[float] = None,
    scales_w: Optional[float] = None,
) -&gt; Tensor:
    return _upsample_nearest(input, output_size, [scales_h, scales_w])
    
@pw_cast_for_opmath
def _upsample_nearest(
    input: Tensor,
    output_size: List[int],
    scales: List[Optional[float]],
    exact: bool = False,
) -&gt; Tensor:
    spatial_indices = _compute_upsample_nearest_indices(
        input, output_size, scales, exact=exact
    )

    indices = [None, None] + spatial_indices
    result = aten._unsafe_index(input, indices)

    if result.ndim == 4:
        # convert output to correct memory format, if necessary
        memory_format = utils.suggest_memory_format(input)

        # following "heuristic: only use channels_last path when it's faster than the contiguous path"
        n_channels = input.shape[1]
        if input.device.type == "cuda" and n_channels &lt; 4:
            memory_format = torch.contiguous_format

        result = result.contiguous(memory_format=memory_format)
    return result
</code></pre>
<h2><a name="attempted-resolution-3" class="anchor" href="#attempted-resolution-3"></a>Attempted Resolution:</h2>
<p>I try to cancel <code>DispatchKey.CompositeImplicitAutograd</code> of <code>upsample_nearest2d</code>, like:</p>
<pre><code class="lang-auto">def disable_implicit_decomposition():
    '''
    Since torch official will implicitly decompose some aten ops,
    disable some ops here to avoid poor performance after decompose.
    '''
    disable_aten_ops = [
        'aten.upsample_nearest1d.vec', 'aten.upsample_nearest1d.default',
        'aten.upsample_nearest2d.vec', 'aten.upsample_nearest2d.default',
        'aten.upsample_nearest3d.vec', 'aten.upsample_nearest3d.default',
    ]

    for op_override in decomposition_table.keys():
        if str(op_override) in disable_aten_ops:
            if DispatchKey.Autograd in op_override.py_kernels:
                op_override.py_kernels.pop(DispatchKey.Autograd)
            if DispatchKey.CompositeImplicitAutograd in op_override.py_kernels:
                op_override.py_kernels.pop(DispatchKey.CompositeImplicitAutograd)
</code></pre>
<p>However, this modification causes the <code>upsample_nearest2d</code> operator to be dispatched to the C++ implementation of <code>upsample_nearest2d</code>. The <code>input.sizes () </code>function causes an error when torch.compile (dynamic=True) is executed.</p>
<pre><code class="lang-auto">Tensor upsample_nearest2d(
    const Tensor&amp; input,
    at::OptionalIntArrayRef output_size,
    std::optional&lt;ArrayRef&lt;double&gt;&gt; scale_factors) {
  auto osize = compute_output_size(input.sizes(), output_size, scale_factors);
  auto scale_h = get_scale_value(scale_factors, 0);
  auto scale_w = get_scale_value(scale_factors, 1);
  return at::upsample_nearest2d(input, osize, scale_h, scale_w);
}
</code></pre>
<p>raise error:</p>
<pre><code class="lang-auto">Cannot call sizes() on tensor with symbolic sizes/strides
</code></pre>
<h2><a name="questions-and-discussion-points-4" class="anchor" href="#questions-and-discussion-points-4"></a>Questions and Discussion Points</h2>
<p>1、Why add <code>DispatchKey.CompositeImplicitAutograd</code> for <code>upsample_nearest2d</code>? The root cause is that you can’t do input.size () in c++ implementation?</p>
<p>2、Why does the <code>upsample_nearest2d</code> function need to call the <code>_upsample_nearest</code>  and decomposes  to  a bunch of small operators ? The problem I encountered in this scenario can also be solved by directly calling <code>torch.ops.aten.upsample_nearest2d</code> in <code>upsample_nearest2d</code>.</p>
<pre><code class="lang-auto">@register_decomposition([aten.upsample_nearest2d.default, aten.upsample_nearest2d.out])
@aten.upsample_nearest2d.default.py_impl(DispatchKey.CompositeImplicitAutograd)
@aten.upsample_nearest2d.default.py_impl(DispatchKey.Autograd)
@out_wrapper(preserve_memory_format=True, exact_dtype=True)
def upsample_nearest2d(
    input: Tensor,
    output_size: List[int],
    scales_h: Optional[float] = None,
    scales_w: Optional[float] = None,
) -&gt; Tensor:
    # return _upsample_nearest(input, output_size, [scales_h, scales_w])
    return torch.ops.aten.upsample_nearest2d(input, output_size, scales_h, scales_w)
</code></pre>
<p>3、Is there a more appropriate way to achieve my goal without splitting the operator during torch.compile?</p>