[tracking] TorchToLinalg Op Support

[vivekkhandelwal1]

Tracking op support for TorchToLinalg lowering. The ops are ordered based on priority (highest -> lowest) in each section.

IMPORTANT: Mark ops you are working on by hovering over it in the list, then clicking the bullseye symbol to its right. This creates an issue for the op and marks them on the list, avoiding duplicate effort.

Contact:

• @vivekkhandelwal1 (Indian time zone)
• @AmosLewis (in the US Pacific time zone)

If you have questions, please check and ask in the discord LLVM/torch-mlir channel.

Instructions on adding new ONNX or torch OP:
https://github.com/llvm/torch-mlir/blob/main/docs/add_ops.md

Please add e2e operator level test(s) to e2eshark test suite for your newly added ops using instructions at:
https://github.com/nod-ai/SHARK-TestSuite/blob/main/e2eshark/README.md

For ONNX Op Support, see #215

• In-Progress (Count: 2):

  • torch.aten.view to linalg #568 @Shukla-Gaurav
  • torch.aten.replication_pad3d #305

• Unimplemented (Count: 26):

  • torch.aten.as_strided (High Priority)
  • torch.aten.randperm #355
  • torch.aten.resize #663
  • torch.aten.resize_as #667
  • torch.aten.max_unpool2d #764
  • torch.aten.lgamma #317
  • boolean indexing ops: AtenNonzeroOp, AtenIndexTensorOp, AtenMaskedSelectOp llvm/torch-mlir#3293
  • torch.aten._linalg_solve_ex #593
  • torch.aten._trilinear
  • torchvision.deform_conv2d
  • torch.aten._fft_r2c
  • torch.aten._fft_c2c
  • torch.aten.unfold
  • torch.aten.normal.Tensor_Tensor
  • torch.aten._sparse_coo_tensor_with_dims_and_tensors
  • torch.aten.rrelu_with_noise #827
  • torch.aten.pixel_unshuffle
  • torch.aten.upsample_trilinear3d
  • torch.aten._sparse_addmm
  • torch.aten._sample_dirichlet
  • torch.aten._cdist_forward
  • torch.aten.linalg_cholesky_ex
  • torch.aten._pdist_forward
  • torch.aten.poisson
  • torch.aten.histc
  • torch.aten.affine_grid_generator

• Completed (Count: 48):

  • torch.aten.max_pool3d_with_indices - Support added by [TorchToLinalg] Support lowering MaxPool3dWithIndices llvm/torch-mlir#3652
  • torch.aten._linalg_slogdet
  • torch.aten.norm.Scalar
  • torch.aten.hann_window.periodic
  • torch.aten.renorm
  • torch.aten.all.dim
  • torch.aten.triu_indices
  • torch.aten.kthvalue
  • torch.aten.avg_pool2d to linalg #643 @AmosLewis
  • torch.aten.tril_indices
  • torch.aten.max_unpool3d #765
  • torch.aten.lerp.Scalar #356
  • torch.aten.logit #290
  • torch.aten.reflection_pad2d #337
  • torch.aten._weight_norm_interface #757
  • torch.aten.pixel_shuffle (completed in #2590)
  • torch.aten.normal_functional #326
  • torch.aten.adaptive_max_pool2d #159
  • torch.aten.replication_pad2d #286
  • torch.aten.glu #772
  • torch.aten.avg_pool3d to linalg #300
  • Implement torch.aten.diag_embed #288
  • torch.aten.acos #338
  • torch.aten.argmin #339
  • torch.aten.min.dim #340
  • torch.aten.exponential #341
  • torch.aten.reflection_pad1d #293
  • torch.aten.prod.dim_int
  • torch.aten.conv_tbc conv1d and conv3d #345
  • torch.aten.isinf
  • torch.aten._adaptive_avg_pool3d #545
  • torch.aten.normal.float_float #343
  • torch.aten.trace #333
  • torch.aten.adaptive_max_pool3d
  • torch.aten.linalg_cross Implement lowering of torch.aten.linalg_cross llvm/torch-mlir#2986
  • torch.aten.diagonal #297
  • torch.aten.cosh #342
  • torch.aten.remainder.Tensor #661
  • torch.aten.prod
  • torch.aten.fmod.Tensor #351
  • aten.floor_divide.Scalar & aten.div.Scalar_mode #618
  • torch.aten.add to linalg #571
  • torch.aten.native_layer_norm to linalg #569
  • torch.aten.squeeze to linalg #572
  • torch.aten.arange.start_step to linalg #573
  • torch.aten.convolution to linalg #642
  • torch.aten.avg_pool3d to linalg #300 @AmosLewis
  • torch.aten.mm to linalg #665

[jinchen62]

@vivekkhandelwal1 According to the triage of iree onnx tests failures, these ops are missing TorchToLinalg lowering. Please add them to the above list.

To reproduce issue:
build venv following https://github.com/nod-ai/SHARK-TestSuite/tree/main/iree_tests#common-venv-setup-with-deps
run iree-compile iree_tests/onnx/node/generated/TEST_NAME/model.mlir -o test.vmfb --iree-hal-target-backends=llvm-cpu --mlir-print-ir-after-all

• torch.aten.cumsum onnx.CumSum #662
  • test_cumsum_1d
  • test_cumsum_1d_exclusive
  • test_cumsum_1d_reverse
  • test_cumsum_1d_reverse_exclusive
  • test_cumsum_2d_axis_0
  • test_cumsum_2d_axis_1
  • test_cumsum_2d_negative_axis
• torch.aten.nonzero Add support for Nonzero TorchToLinalg lowering #465
  • test_nonzero_example

[lingzhiz1998]

Hi, can I take torch.aten.max_pool3d_with_indices op task?

[lingzhiz1998]

Attach this PR llvm/torch-mlir#3652 to mark max_pool3d_with_indices as in progress. I can not change the status of task list.

[vivekkhandelwal1]

Attach this PR llvm/torch-mlir#3652 to mark max_pool3d_with_indices as in progress. I can not change the status of task list.

Done