Supporting the Core ATen Opset

[wonjoolee95]

🚀 Feature

As PyTorch/XLA tries to support the PyTorch core ATen opset in its entirety, it requires lowering each core ATen op in PyTorch/XLA. As of today, most of these ops are already lowered. This issue will serve as a tracker for our support of the core ATen opset.

Unlowered core ATen ops

• AtenAcos_AutoI32
• AtenAcosh_AutoI32
• AtenArangeStartStep
• AtenAsStrided
• AtenAsStrided_AutoI32
• AtenAsin_AutoI32
• AtenAsinh_AutoI32
• AtenAtan2_AutoI32
• AtenAtan_AutoI32
• AtenAvgPool2dBackward
• AtenAvgPool2dBackward_AutoI32
• AtenCdistForward_AutoI32
• AtenCol2im
• AtenConvolutionBackward
• AtenCos_AutoI32
• AtenCosh_AutoI32
• AtenEmbeddingDenseBackward
• AtenEmptyStrided
• AtenErf_AutoI32
• AtenExp_AutoI32
• AtenExpm1_AutoI32
• AtenFull
• AtenGelu_AutoI32
• AtenGridSampler2d
• AtenGridSampler2d_AutoI32
• AtenIsinf_AutoI32
• AtenIsnan_AutoI32
• AtenMaxPool2dWithIndicesBackward
• AtenMeanDim_AutoI32
• AtenMean_AutoI32
• AtenNativeBatchNormLegitNoStats
• AtenNativeBatchNormLegitNoStats_AutoI32
• AtenNativeBatchNormLegit
• AtenNativeBatchNormLegit_AutoI32
• AtenNativeGroupNormBackward
• AtenNativeGroupNorm_AutoI32
• AtenNativeLayerNormBackward
• AtenNativeLayerNorm_AutoI32
• AtenPdistForward
• AtenPdistForward_AutoI32
• AtenPixelShuffle
• AtenPixelShuffle_AutoI32
• AtenPowTensorTensor_AutoI32
• AtenRand
• AtenRandn
• AtenRandperm_AutoI32
• AtenReflectionPad3d
• AtenReflectionPad3d_AutoI32
• AtenReplicationPad3d
• AtenReplicationPad3d_AutoI32
• AtenRsqrt_AutoI32
• AtenScalarTensor
• AtenSigmoid_AutoI32
• AtenSin_AutoI32
• AtenSliceScatter
• AtenSqrt_AutoI32
• AtenSymNumel
• AtenSymSizeInt
• AtenSymStorageOffset
• AtenSymStrideInt
• AtenTopk

[danielvegamyhre]

Would working on lowering some of these core ATen ops into Pytorch/XLA be a good first issue for a new contributor? If so, I can work on some, perhaps starting with AtenAcos_AutoI32?

[JackCaoG]

Absolutely! You can take a look at https://github.com/pytorch/xla/blob/master/OP_LOWERING_GUIDE.md and https://github.com/pytorch/xla/blob/master/CODEGEN_MIGRATION_GUIDE.md to get some background on how to do lowering. You can also check recently merged pr of lowering for reference.

[danielvegamyhre]

Awesome, I'll take a look, thanks!

[miladm]

cc @qihqi
can you please identify which ops are covered? cc @wonjoolee95 to assist with closing the gap with you.

[qihqi]

Hi,

I created a unit test with failed ops that failed conversion to StableHLO: https://github.com/pytorch/xla/pull/5796/files

To add support an op the process would be roughly:

1. find a test marked with unittest.expectFailure or unittest.skip. Remove that annotation
2. Run this test and see the failure.
3. Fix the test. Sometimes the fix is to modify the test itself because the input to the test is automatically generated and might not be valid input to begin with.

[wonjoolee95]

Marking this as complete, https://github.com/pytorch/xla/projects?type=classic&query=Core+Aten+Opset.