TorchInductor missing ops tracker

[jansel]

The following ops are using ir.FallbackKernel via make_fallback() in lowering.py and appear in benchmarks. We should rewrite them to use decomps or lowerings.

• Add decomp/lowering: aten.as_strided_scatter #93650
• aten.grid_sampler_2d_backward (higher priority)
• aten.upsample_bilinear2d_backward (higher priority)
• aten._adaptive_avg_pool2d_backward
• aten.upsample_bicubic2d_backward
• aten._fused_moving_avg_obs_fq_helper
• aten.upsample_nearest3d (needed for FAIR model)
• aten.avg_pool3d (needed for FAIR model)
• aten.bucketize (needed for internal model) - not targeting for codegen, do a fallback
• aten.prod (needed for research model) - not targeting for codegen, do a fallback

Might not be possible (in a performant way), but currently use fallbacks:

• aten.convolution_backward (might need to hold of on this if perf doesn't match)
• aten._cudnn_rnn (might need to hold of on this if perf doesn't match) - not targeting for codegen, do a fallback
• aten._cudnn_rnn_backward (might need to hold of on this if perf doesn't match) - not targeting for codegen, do a fallback
• aten._embedding_bag (may have a template internally) (Attempted with [No CI] Decomp for _embedding_bag #84235, but it's very hard to make it performant and inductor-friendly)
• [inductor] Lower aten.cumsum #93631 not targeting for codegen, do a fallback
• torchvision.roi_align (need to sort of decomps for domain libs)

Done:

• aten._adaptive_avg_pool2d
• aten.nll_loss_forward (PR)
• aten.grid_sampler_2d (could use the version from here)
• aten.upsample_bilinear2d Add decomposition for aten.upsample_bilinear2d.vec #80964
• aten.reflection_pad2d_backward (higher priority)
• aten.native_batch_norm_backward (move aten.native_batch_norm_backward decomposition to core #81522)
• aten.avg_pool2d_backward
• aten.expand_as
• aten.glu_backward (Added glu_backward decomp #78919)
• aten.masked_fill_
• aten.max_pool2d_with_indices_backward
• aten.select_scatter
• aten.upsample_nearest2d_backward
• aten.baddbmm
• aten.log1p
• aten.bernoulli_
• aten.conj_physical
• [inductor] Lower aten.index_add_/aten.index_add torchdynamo#885
• aten.im2col_backward
• aten.native_group_norm_backward ([Prim] Implement group_norm_backward #84037)
• aten.im2col (Decomp for aten.im2col #84303)
• aten.binary_cross_entropy_with_logits (needed for internal model)
• aten.upsample_bicubic2d (adds nn.Crop1d, nn.Crop2d, nn.Crop3d, see #1331 #1349)
• aten.unfold Add refs for torch.unfold and a decomposition for its backward. #85629
• aten.unfold_backward Add refs for torch.unfold and a decomposition for its backward. #85629
• aten.col2im Implement col2im decomposition and fix im2col and add a few preconditions #85541
• aten.mse_loss_backward (higher priority)
• aten.softplus (needed for research model)
• aten.softplus_backward (needed for research model)

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @soumith @ngimel

[Chillee]

aten.convolution_backward

Unfortunately, there's no good way of decomposing this today. The gradient (wrt the weights iirc?) cannot be expressed with a forwards convolution.

aten.nll_loss_forward

This one is WIP: #78491

[ezyang]

Unfortunately, there's no good way of decomposing this today. The gradient (wrt the weights iirc?) cannot be expressed with a forwards convolution.

Well, to be more precise, if we have a suitably generalized convolution that has output shape and transposition that would be sufficient for all convolutions forwards/backwards.

[soumith]

The gradient (wrt the weights iirc?) cannot be expressed with a forwards convolution.

It can. Where you have to pad the output image appropriately, and the convolve it with a transposed kernel weight matrix.

See the last diagram in http://soumith.ch/ex/pages/2014/08/07/why-rotate-weights-convolution-gradient/

[Chillee]

aten.log1p

btw, this is currently a prim in Primtorch. So there won't be a further decomposition for it.

aten.expand_as

Also, this is a CompositeImplicitAutograd op - where are you seeing this today?

[jansel]

aten.expand_as

Also, this is a CompositeImplicitAutograd op - where are you seeing this today?

I am seeing it in AOT Autograd graphs for training. I see it in multiple models for example BERT_pytorch.

[Chillee]

Interesting, will have to check where it comes from. Perhaps it's in a decomposition and we're not properly recursively decomposing? But I don't see it being used in any decomposition right now besides embedding_dense_backward.

[soumith]

Well, to be more precise, if we have a suitably generalized convolution that has output shape and transposition that would be sufficient for all convolutions forwards/backwards.

A convolution that has strides, dilation and zero padding can express all needs in a closed form.
Convolution forward, backward-input, backward-weight, etc. can be expressed in terms of this convolution.

[lezcano]

I'll have a stab at aten._adaptive_avg_pool2d and its backward.

[jansel]

I'll have a stab at aten._adaptive_avg_pool2d and its backward.

@lezcano I think @eellison just did that one here https://github.com/pytorch/torchdynamo/pull/962/files

It requires ops.masked so it had to be done as a lowering rather than a decomp.

[lezcano]

AGGGGG. Would it be possible to update the OP with the ops that have already been implemented?

Also, I was able to implement this one without raw loops (but still peeking twice into the data, one per dimension). Do you think this could be of interest by itself? I am writing this in _decomp, I am not sure what's the relation between _decomp and these lowerings. Which one is preferable? The lowerings use some constructions I don't understand at first sight.