Push for finishing transition to Python dtype functions

[ramiro050]

Hi everyone,

I would like to get some help from the community to finish up the transition from using RefineTypes.cpp for encoding how to calculate the resulting dtypes of the tensors of all the ops to using dtype functions that are written in Python. For more information about the rationale behind this transition see the Custom Op RFC.

The transition will take place on the branch dtype-functions-staging, so that it happens in a single go in main, otherwise many large workloads can fail to lower through Torch-MLIR. The reason things can fail is that if you have some ops being handled by RefineTypes and other ops using the Python dtype functions that get applied in the TorchDtypeRefinementPipeline, then the RefineTypes pass and the TorchDtypeRefinementPipeline will cycle over and over, each pass propagating dtype information until encountering an op that is not handled by the current pass. For workloads with many ops, the cycling of passes can easily reach the iteration limit in the LowerToBackendContract pass, resulting in a failure when lowering through Torch-MLIR. Currently, there are only a few ops using dtype functions in main, and it is already affecting some user workloads.

To avoid breaking workloads for users, or forcing users to increase the iteration limit affecting performance of Torch-MLIR, we will instead do the entire transition on the branch dtype-functions-staging, where we can XFAIL complex tests temporarily as needed.

The contribution process

1. Choose an item from the work list below and edit this issue by adding your name alongside the checklist item you're working on to avoid duplicate work
   • If you don't have edit access, just mention it as a comment, and I will update the issue
2. Make a new branch off of the dtype-functions-staging branch
3. Remove the op or ops linked in the checklist item from RefineTypes.cpp
4. Add dtypes functions with proper testing to the abstract_interp_lib_gen.py. See the docs for information on how two write a dtype function
5. Run ./build_tools/update_abstract_interp_lib.sh to update abstract interp library
6. Your patch is only expected to pass the e2e tests that test a single op. If your patch causes tests that use many ops like ResNet18 or Mobilenet to fail, mark those tests as XFAIL
7. Make a PR to merge your branch into dtype-functions-staging (see the PR for ops that return i1s for reference)
8. Tag me (@ramiro050) for review
9. Once merged, mark the item in the work list as finished

Work List

• dtype from first operand (@li-plus Add dtype functions for ops that take dtype from 1st operand #1895)
• dtype is always float32 (@li-plus Add dtype functions for floating point ops #1813)
• dtype from second operand (@li-plus Add dtype functions for ops that take dtype from 2nd operand #1891)
• dtype is always i1 (@ramiro050 Add dtype functions for comparison ops #1806)
• dtype is always si64 (@ramiro050 Add more dtype funcs #1962)
• promote two dtypes assuming non-zero rank (@ramiro050 Add dtype functions for two tensor promotion ops #1831 Change dtype functions interface to take ints tuple for each tensor #1865)
• promote two dtypes assuming possibly zero rank (@ramiro050 Add dtype functions for two tensor promotion ops #1831)
• promote three dtypes (@ramiro050 Add more dtype funcs #1962)
• promote LHS with scalar RHS (@ramiro050 Add more dtype funcs #1962)
• promote 2nd and 3rd operands (@ramiro050 Add more dtype funcs #1962)
• 2 results take dtype from first operand (@ramiro050 Add dtype functions for "2/3 results take dtype from first operand" #1967)
• 3 results take dtype from first operand (@ramiro050 Add dtype functions for "2/3 results take dtype from first operand" #1967)
• Arange ops (@ramiro050 Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057)
• Reduction ops and another reduction op (@ramiro050 Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057)
• Scalar to tensor ops (@ramiro050 Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057)
• Constant tensor allocation ops (@ramiro050 Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057)
• Type conversion ops (@ramiro050 Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057)
• Embedding ops (@ramiro050 Replace RefineTypes with dtype functions #2105)
• Softmax ops (@ramiro050 Replace RefineTypes with dtype functions #2105)
• Binary scalar ops(@ramiro050 Replace RefineTypes with dtype functions #2105)
• Rand ops and more rand ops (@ramiro050 Add dtype functions for rand* and var_mean* ops #2089)
• VarMean ops (@ramiro050 Add dtype functions for rand* and var_mean* ops #2089)
• Miscellaneous ops: AtenAtan2Op, AtenLinearOp, AtenMaxPool2dWithIndicesOp, AtenTensorOp, AtenCatOp, Aten_ShapeAsTensorOp, AtenScalarImplicitOp, PrimNumToTensorScalarOp (@ramiro050 Add miscellaneous dtype functions #2090)

Happy to help with any questions/issues!

[li-plus]

I'm willing to help. Could I take the the first 3 tasks?

[ramiro050]

I'm willing to help. Could I take the the first 3 tasks?

Awesome! I've updated the work list

[powderluv]

@gpetters94 can you please help chip away at these ops too.

[gpetters94]

@ramiro050 I can take Embedding and Softmax for now.