Skip to content

Add information about codegen in codegen/xla_native_functions.yaml#8817

Merged
pgmoka merged 8 commits intomasterfrom
add_info_on_codegen_settings
Mar 17, 2025
Merged

Add information about codegen in codegen/xla_native_functions.yaml#8817
pgmoka merged 8 commits intomasterfrom
add_info_on_codegen_settings

Conversation

@pgmoka
Copy link
Copy Markdown
Collaborator

@pgmoka pgmoka commented Mar 11, 2025

Add some information on codegen and supported configs based on #8713.

Info on generation also comes from https://github.com/pytorch/pytorch/blob/main/torchgen/gen_lazy_tensor.py.

@pgmoka pgmoka requested review from tengyifei and ysiraichi March 11, 2025 20:26
@pgmoka pgmoka self-assigned this Mar 11, 2025
@pgmoka
Copy link
Copy Markdown
Collaborator Author

pgmoka commented Mar 11, 2025

@tengyifei @ysiraichi my knowledge on the other forms of generation (symint, and autograd) is a little spotty. I would be interested if y'all have anything to add here.

Copy link
Copy Markdown
Collaborator

@ysiraichi ysiraichi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks nice. Besides the comments, I have 2 suggestions:

  • Add symint and autograd descriptions
  • Mention what gets generated in which file

Let me know if you need help on gathering this information.

Comment thread codegen/xla_native_functions.yaml Outdated
Comment thread codegen/xla_native_functions.yaml Outdated
@pgmoka
Copy link
Copy Markdown
Collaborator Author

pgmoka commented Mar 12, 2025

This looks nice. Besides the comments, I have 2 suggestions:

  • Add symint and autograd descriptions
  • Mention what gets generated in which file

Let me know if you need help on gathering this information.

@ysiraichi I haven't been able to find a great definite source on these which is why I haven't written about them. If you have links, I would appreciate it.

Comment thread codegen/xla_native_functions.yaml Outdated
Comment thread codegen/xla_native_functions.yaml
Comment thread codegen/xla_native_functions.yaml Outdated
Comment thread codegen/xla_native_functions.yaml Outdated
- triu
- trunc
# ir_gen: Operations in ir_gen generate only the declaration of their respective
# lazy IR node classes. Their native functions are not generated.
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Double checking my understanding, this means ir_gen doesn't generate native functions, but supported does generate native functions? Or is there a difference between native function declaration and native function definition (i.e. implementation)?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this means ir_gen doesn't generate native functions, but supported does generate native functions?

Yes.

Or is there a difference between native function declaration and native function definition?

ir_gen generates the lazy IR node class declaration. Some of the class' functions still have to be manually written (e.g. XXX::Lower()).

supported generates the XLANativeFunction::XXX declaration + registration to the XLA dispatch-key. We still have to manually write the definition of that function.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there any recommendations of something to add here from this conversation?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's include Yukio's explanation into the yaml docs

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have added the explanations from this comment

Comment thread codegen/xla_native_functions.yaml Outdated
@ysiraichi
Copy link
Copy Markdown
Collaborator

Here's a small summary about each component inside xla_native_functions.yaml file:

  • full_codegen
    • Generates:
      • Native function declaration and definition (at XLANativeFunctions.{h,cpp})
      • IR node class declaration (at LazyIR.h)
      • Kernel registration (at RegisterXLA.cpp)
    • Needs:
      • <Op>OutputShape() declaration and definition (at ops_xla_shape_fn.{h,cpp})
      • <Op>::Lower() definition (at ops_lower_fn.cpp)
  • supported
    • Generates:
      • Native function declaration (at XLANativeFunctions.h)
      • Kernel registration (at RegisterXLA.cpp)
    • Needs:
      • Native function definition (at aten_xla_type.cpp)
  • ir_gen (at the moment, always used with supported)
    • Generates:
      • IR node class declaration (at LazyIR.h)
  • symint (at the moment, always used with supported)
    • Same as supported (or full_codegen, if used instead), but uses the operation overload that takes SymInts instead of plain integers
  • autograd
    • Generates:
      • Native function declaration (at XLANativeFunctions.h)
      • Kernel registration (at RegisterAutogradXLA.cpp)
    • Needs:
      • Native function definition (at aten_xla_type.cpp)

Note: this is my own understanding by looking at the source code.

@pgmoka
Copy link
Copy Markdown
Collaborator Author

pgmoka commented Mar 14, 2025

  • symint

I have added a description to symint based on this and #3978. Please let me know what you think of it @ysiraichi

@pgmoka pgmoka requested a review from tengyifei March 14, 2025 18:43
Copy link
Copy Markdown
Collaborator

@ysiraichi ysiraichi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be nice to write down what's getting generated in which files. And, as @tengyifei mentioned, what's needed for finishing the lowering.

Comment thread codegen/xla_native_functions.yaml Outdated
Comment thread codegen/xla_native_functions.yaml Outdated
Copy link
Copy Markdown
Collaborator

@tengyifei tengyifei left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM modulo the open comments

@pgmoka
Copy link
Copy Markdown
Collaborator Author

pgmoka commented Mar 15, 2025

I think it would be nice to write down what's getting generated in which files. And, as @tengyifei mentioned, what's needed for finishing the lowering.

I think this information should be in https://github.com/pytorch/xla/blob/master/docs/source/contribute/op_lowering.md and https://github.com/pytorch/xla/blob/master/docs/source/contribute/codegen_migration.md. I would recommend link to those documents, and then supplementing them if we believe there are gaps. We could create an issue to add more information to https://github.com/pytorch/xla/blob/master/docs/source/contribute/op_lowering.md on symint and autograd, but I would do this in a follow-up PR. Let me know what you think @ysiraichi.

@pgmoka
Copy link
Copy Markdown
Collaborator Author

pgmoka commented Mar 15, 2025

I have added links to the referred docs in the head of the file.

@pgmoka pgmoka requested a review from ysiraichi March 15, 2025 01:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants