Push for Bert, MaskRCNN and Resnet-18 E2E support using Lazy Tensor Core

[ramiro050]

The purpose of this bug is to help keep track of tasks needed to have E2E support for Bert, MaskRCNN, and Resnet-18 using the Lazy Tensor Core (LTC) and lowering through torch-mlir.

At the moment, the main part that is missing is op support in LTC. Below is a table of the LTC ops needed. The list of ops was determined by running two scripts: this MaskRCNN script which generates this output, and this Resnet script which generates this output. The ops missing are the ones with the aten:: prefix in the output of the scripts. For more information of how to setup and run the examples, see here.

LTC Ops Needed

Status column symbols:
- unimplemented
+ started
x finished

Ops	Status	Owner	Model	Notes
aten::_index_put_impl_	-		MaskRCNN
aten::arange.start_out	+	silvasean	MaskRCNN
aten::exp.out	x	ramiro050	MaskRCNN	pytorch/pytorch#67213
aten::floor.out	x	ramiro050	MaskRCNN	pytorch/pytorch#66770
aten::index.Tensor	+	ramiro050	MaskRCNN
aten::log2.out	x	ramiro050	MaskRCNN	pytorch/pytorch#66771
aten::max_pool2d_with_indices	+	vivekkhandelwal1	MaskRCNN, Resnet
aten::upsample_nearest2d.out	-		MaskRCNN
aten::mean.out	x	alanwaketan	Resnet	pytorch/pytorch#67174
aten::sort	x	silvasean	Resnet	pytorch/pytorch#67053

torch-mlir ops Needed

Below is a list of ops needed on the torch-mlir side. This list was compiled by going over the ops detected by LTC when running this MaskRCNN script (output with list of ops detected can be found here), this Bert script (output), and the Resnet-18 model in the PyTorch benchmarks (instructions for setting it up with LTC), and checking which had lowerings in torch-mlir and which did not.

Note: Bert and Resnet18 are currently the only training models. The ops needed for MaskRCNN training will be added soon.

Status column symbols:
- unimplemented
+ started
x finished

The full op lists including finished ones are moved to #365 (comment). This new table only contains ops to be done so that we can be more focused.

Op	Status	Owner	Model	Notes
aten::bernoulli_	+	pashu123	Bert Training	rng op
aten::embedding_dense_backward	+	vivekkhandelwal1	Bert Training	histogram
aten::native_layer_norm_backward	+	gprateek93	Bert Training	PR546, PR570
aten::nll_loss_backward	+	pashu123	Bert Training, Resnet-18 Training	PR463
aten::_copy_from	+	pashu123	Resnet-18 Training, MaskRCNN Inference	torchscript baseline won't run (to be investigated)
aten::convolution_backward_overrideable	+	gpetters94	Resnet-18 Training
aten::max_pool2d_with_indices_backward	+	vivekkhandelwal1	Resnet-18 Training
aten::native_batch_norm_backward	+	Shukla-Gaurav	Resnet-18 Training
aten::native_batch_norm	+	Shukla-Gaurav	Resnet-18 Training	PR563
aten::random_.to	+	gprateek93	Resnet-18 Training	rng op
aten::_copy_from_and_resize	+	gpetters94	Resnet-18 Training, MaskRCNN Inference
aten::convolution_overrideable	+	gpetters94	Resnet-18 Training, MaskRCNN Inference
aten::convolution	+	gpetters94	Resnet-18 Training(through AOTAutograd)
aten::convolution_backward	+	gpetters94	Resnet-18 Training(through AOTAutograd)
aten::max_pool2d_with_indices	+	vivekkhandelwal1	Resnet-18 Training, MaskRCNN Inference	#518
aten::_index_put_impl_	-		MaskRCNN Inference	histogram
aten::stack	-	pashu123	MaskRCNN Inference
aten::topk	-	gprateek93	MaskRCNN Inference
aten::upsample_nearest2d	-	gprateek93	MaskRCNN Inference
torchvision::nms	-		MaskRCNN Inference
torchvision::roi_align	-		MaskRCNN Inference

[silvasean]

Awesome. Just as an FYI, I'm working on plugging torch-mlir into LTC.

[Eliasj42]

@silvasean Hi, I just started working with the nod team. Is there any way I could get editing permissions for this issue?

Thanks

[silvasean]

@silvasean Hi, I just started working with the nod team. Is there any way I could get editing permissions for this issue?

Thanks

Done

[vivekkhandelwal1]

@silvasean Hi, I have started working with the Nod AI; please add me to this issue.

[silvasean]

Done!

[cathyzhyi]

@Eliasj42 could you update the ops status table with the PR numbers?

[ramiro050]

If you encounter an op that has many different variants, you can search here https://gist.github.com/ramiro050/349b61aa93a0b136bfc5c363a6a21398 to see the signature of the op being used by each of the three models

[Nirvedh]

Hi, I recently started working at Nod.ai and I am working on aten::t

[cathyzhyi]

Original op list table

Op	Status	Owner	Model	Notes
aten::bmm	x		Bert Training
aten::gelu	x		Bert Training
aten::tanh	x		Bert Training
aten::transpose	x		Bert Training
aten::_softmax_backward_data	X	cathyzhyi	Bert Training	PR402
aten::bernoulli_	-		Bert Training	hold: RNG ops
aten::embedding_dense_backward	+	vivekkhandelwal1	Bert Training
aten::gelu_backward	x	gprateek93	Bert Training	PR418
aten::index_select	x	vivekkhandelwal1	Bert Training	#466
aten::native_layer_norm_backward	+	gprateek93	Bert Training
aten::native_layer_norm	x	gprateek93	Bert Training	PR470
aten::tanh_backward	x	pashu123	Bert Training	PR408
aten::permute	x		Bert Training, MaskRCNN Inference
aten::sub	x		Bert Training, MaskRCNN Inference
aten::unsqueeze	x		Bert Training, MaskRCNN Inference
aten::_softmax	x	gprateek93	Bert Training, MaskRCNN Inference	PR431
aten::select (select.int)	X		Bert Training, MaskRCNN Inference
aten::slice(slice.Tensor)	X		Bert Training, MaskRCNN Inference
aten::mm	x		Bert Training, Resnet-18 Training
aten::sum	x		Bert Training, Resnet-18 Training
aten::_log_softmax_backward_data	x	pashu123	Bert Training, Resnet-18 Training	PR413
aten::_log_softmax	x	pashu123	Bert Training, Resnet-18 Training	PR397
aten::addcdiv	X	Eliasj42	Bert Training, Resnet-18 Training	PR442
aten::addcmul	X	Eliasj42	Bert Training, Resnet-18 Training	PR442
aten::nll_loss_backward	+	pashu123	Bert Training, Resnet-18 Training	[PR463]
aten::nll_loss_forward	X	pashu123	Bert Training, Resnet-18 Training	[PR452]
aten::add	x		Bert Training, Resnet-18 Training, MaskRCNN Inference
aten::div	x		Bert Training, Resnet-18 Training, MaskRCNN Inference
aten::mul	x		Bert Training, Resnet-18 Training, MaskRCNN Inference
aten::view	x		Bert Training, Resnet-18 Training, MaskRCNN Inference
aten::_copy_from	+	pashu123	Bert Training, Resnet-18 Training, MaskRCNN Inference	torchscript baseline won't run (to be investigated)
aten::addmm	X	Shukla-Gaurav	Bert Training, Resnet-18 Training, MaskRCNN Inference	PR412
aten::expand	X	Shukla-Gaurav	Bert Training, Resnet-18 Training, MaskRCNN Inference	PR396
aten::fill_	x	pashu123	Bert Training, Resnet-18 Training, MaskRCNN Inference	PR440
aten::sqrt	x		Bert Training, Resnet-18 Training, MaskRCNN Inference
aten::t	x	nirvedhmeshram	Bert Training, Resnet-18 Training, MaskRCNN Inference	PR492
aten::zero_	x	vivekkhandelwal1	Bert Training, Resnet-18 Training, MaskRCNN Inference	#464
aten::convolution_backward_overrideable	-		Resnet-18 Training
aten::max_pool2d_with_indices_backward	-		Resnet-18 Training
aten::mean	x	vivekkhandelwal1	Resnet-18 Training	#430
aten::native_batch_norm_backward	-	Shukla-Gaurav	Resnet-18 Training
aten::native_batch_norm	+	Shukla-Gaurav	Resnet-18 Training
aten::random_.to	-		Resnet-18 Training
aten::threshold_backward	x	vivekkhandelwal1	Resnet-18 Training	#509
aten::relu	x		Resnet-18 Training, MaskRCNN Inference
aten::_copy_from_and_resize			Resnet-18 Training, MaskRCNN Inference
aten::convolution_overrideable	-		Resnet-18 Training, MaskRCNN Inference
aten::max_pool2d_with_indices	+	vivekkhandelwal1	Resnet-18 Training, MaskRCNN Inference	#518
aten::cat	x		MaskRCNN Inference
aten::exp	x		MaskRCNN Inference
aten::sigmoid	x		MaskRCNN Inference
aten::clamp	x		MaskRCNN Inference
aten::_index_put_impl_	-		MaskRCNN Inference
aten::arange.start	x	vivekkhandelwal1	MaskRCNN Inference	#501
aten::bitwise_and	x	vivekkhandelwal1	MaskRCNN Inference	#435
aten::eq	x	gpetters94	MaskRCNN Inference
aten::floor	x		MaskRCNN Inference
aten::ge	x	gpetters94	MaskRCNN Inference
aten::gt	+	gpetters94	MaskRCNN Inference
aten::index.Tensor	+	vivekkhandelwal1	MaskRCNN Inference	#495
aten::log2	x	gpetters94	MaskRCNN Inference
aten::lt	x	gpetters94	MaskRCNN Inference
aten::nonzero	x	gpetters94	MaskRCNN Inference
aten::rsqrt	x	gpetters94	MaskRCNN Inference
aten::squeeze	x	Shukla-Gaurav	MaskRCNN Inference	PR441 PR453
aten::stack	-	pashu123	MaskRCNN Inference
aten::topk	-	gprateek93	MaskRCNN Inference
aten::upsample_nearest2d	-	gprateek93	MaskRCNN Inference
torchvision::nms	-		MaskRCNN Inference
torchvision::roi_align	-		MaskRCNN Inference

[powderluv]

should we trim the list to only what needs to be added will be easier to track

[cathyzhyi]

should we trim the list to only what needs to be added will be easier to track

yes, I put the trimmed table in the original description of the issue.