Skip to content

Inline all scalars by default in export path#6803

Merged
lsy323 merged 3 commits intomasterfrom
lsiyuan/do-no-lift-constant
Mar 22, 2024
Merged

Inline all scalars by default in export path#6803
lsy323 merged 3 commits intomasterfrom
lsiyuan/do-no-lift-constant

Conversation

@lsy323
Copy link
Copy Markdown
Collaborator

@lsy323 lsy323 commented Mar 22, 2024

Motivation of the change:

Currently only special scalars (0 or 1) will be inlined in the lowered HLO/StableHLO. In export path with HLFB, lifted scalars generated during tracing would result in additional arg in the stablehlo.composite. This makes downstream compiler confusing when mapping the args in the model code and stablehlo.composite.

Inlining scalars should work well in the export path. Because torch.export doesn't trace non-tensor input at this moment, all non-tensor values are burned in the exported FX graph (with dangling placeholders).

Change:

  • Expose a Python API to config Lazy Tensor Core to inline all at::Scalar.
  • Add an option in StableHLO export to allow caller to config to inline all scalars, or only inline scalars with special values. By default, export will inline all scalars.

Test:
Added new tests cases.

@lsy323 lsy323 requested review from JackCaoG and qihqi March 22, 2024 00:47
@lsy323 lsy323 changed the title Inline all scalar by default in export Inline all scalars by default in export path Mar 22, 2024
@lsy323 lsy323 merged commit f43ac41 into master Mar 22, 2024
@lsy323 lsy323 deleted the lsiyuan/do-no-lift-constant branch March 22, 2024 21:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants