feat: add support for custom compile options in torch_xla.compile and… #9575

sshonTT · 2025-08-21T20:05:49Z

… PJRT backend

Issue description : #9555

This change introduces the ability to pass custom compile options from Python down to the PJRT backend, allowing users to fine-tune XLA compilation behavior without modifying core code. Key changes:

Python API
- Added custom_compile_options parameter to torch_xla.compile for passing compile-time options as a dict (supports bool, float, int, and str values).
- Added torch_xla.set_custom_compile_options() utility for setting compile options globally.
- Added internal binding _XLAC._set_custom_compile_options().
C++ Runtime
- Added SetCustomCompileOptions() virtual method to ComputationClient and implemented it in PjRtComputationClient.
- PjRtComputationClient now stores custom_compile_options_ and injects them into xla::CompileOptions.env_option_overrides during compilation.
- Options are stringified before being passed to XLA for compatibility. Motivation: This enables advanced users to pass through backend-specific tuning flags (e.g., enabling experimental optimizations, toggling partitioning strategies) without hardcoding them, improving flexibility for research and debugging workflows.

sshonTT · 2025-08-21T20:27:45Z

@qihqi Thank you for your review! I realized there was a lint issue and have fixed it. Would you mind re-running the tests?

qihqi · 2025-08-25T16:39:34Z

hi @sshonTT would you rebase to latest HEAD? the 2 CI issue should be fixed

sshonTT · 2025-08-26T01:03:42Z

hi @sshonTT would you rebase to latest HEAD? the 2 CI issue should be fixed

@qihqi thank you for the info. I had rebased so could you re-run test?

sshonTT · 2025-08-29T12:47:59Z

hi @qihqi , gentle reminder.

sshonTT · 2025-09-02T15:34:19Z

Hi @zhanyong-wan — not sure if this is the right tag. I’m waiting on @qihqi’s follow-up (he previously approved). I rebased per his suggestion and want to confirm CI is green. Could you trigger CI for this PR and, if you have time, give it a review? Thanks!

zhanyong-wan

Thanks!

torch_xla/csrc/runtime/computation_client.h

torch_xla/torch_xla.py

torch_xla/csrc/runtime/computation_client.h

zhanyong-wan

Thanks!

torch_xla/torch_xla.py

zhanyong-wan

In the future, could you avoid git rebase or git merge during a code review? Otherwise it's hard to track the changes. Thanks!

torch_xla/torch_xla.py

torch_xla/csrc/runtime/pjrt_computation_client.cpp

sshonTT · 2025-09-11T20:11:42Z

hi @qihqi, I see #9634 fixes build failure in ci, so I rebased it.

sshonTT · 2025-10-02T14:28:22Z

I noticed the efforts in #9650 and #9655 regarding the current build failure. Hopefully they will resolve the issue. I’ve rebased once again.

… PJRT backend This change introduces the ability to pass custom compile options from Python down to the PJRT backend, allowing users to fine-tune XLA compilation behavior without modifying core code. Key changes: * Python API * Added custom_compile_options parameter to torch_xla.compile for passing compile-time options as a dict (supports bool, float, int, and str values). * Added torch_xla.set_custom_compile_options() utility for setting compile options globally. * Added internal binding _XLAC._set_custom_compile_options(). * C++ Runtime * Added SetCustomCompileOptions() virtual method to ComputationClient and implemented it in PjRtComputationClient. * PjRtComputationClient now stores custom_compile_options_ and injects them into xla::CompileOptions.env_option_overrides during compilation. * Options are stringified before being passed to XLA for compatibility. Motivation: This enables advanced users to pass through backend-specific tuning flags (e.g., enabling experimental optimizations, toggling partitioning strategies) without hardcoding them, improving flexibility for research and debugging workflows.

sshonTT · 2025-10-02T18:30:42Z

I just noticed that GetComputationClientOrDie API is deprecated after https://github.com/pytorch/xla/pull/9620/files which is my fault, had to care about it. I just remove GetComputationClientOrDie.

sshonTT · 2025-10-07T17:12:47Z

@zhanyong-wan , @qihqi Could you re-trigger this test? I had checked on my local, and want to see if it is working.

sshonTT · 2025-10-09T13:20:38Z

Thanks for running CI! Can this be merged in?

zhanyong-wan · 2025-10-09T15:07:45Z

Yes, please resolve the comments and merge.

qihqi approved these changes Aug 21, 2025

View reviewed changes

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from b3e8b3e to 2a4b284 Compare August 21, 2025 20:25

qihqi enabled auto-merge (squash) August 23, 2025 04:04

auto-merge was automatically disabled August 26, 2025 01:02
Head branch was pushed to by a user without write access

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from 2a4b284 to 2a2eed5 Compare August 26, 2025 01:02

zhanyong-wan requested changes Sep 2, 2025

View reviewed changes

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch 3 times, most recently from 2bdd9ed to 7be5b39 Compare September 2, 2025 19:21

zhanyong-wan reviewed Sep 3, 2025

View reviewed changes

torch_xla/torch_xla.py Outdated Show resolved Hide resolved

torch_xla/torch_xla.py Outdated Show resolved Hide resolved

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from 7be5b39 to 4278f97 Compare September 4, 2025 18:22

zhanyong-wan reviewed Sep 5, 2025

View reviewed changes

torch_xla/torch_xla.py Outdated Show resolved Hide resolved

torch_xla/torch_xla.py Outdated Show resolved Hide resolved

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from 4278f97 to f3ad8bd Compare September 5, 2025 15:59

zhanyong-wan approved these changes Sep 8, 2025

View reviewed changes

torch_xla/csrc/runtime/pjrt_computation_client.cpp Outdated Show resolved Hide resolved

qihqi enabled auto-merge (squash) September 9, 2025 01:12

auto-merge was automatically disabled September 11, 2025 20:10
Head branch was pushed to by a user without write access

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from f3ad8bd to a27f105 Compare September 11, 2025 20:10

zhanyong-wan approved these changes Sep 16, 2025

View reviewed changes

sshonTT mentioned this pull request Sep 19, 2025

Add support for parsing custom compiler options in PJRT backend compile tenstorrent/tt-xla#1321

Merged

1 task

qihqi enabled auto-merge (squash) September 19, 2025 17:00

sshonTT mentioned this pull request Sep 24, 2025

Fix Terraform usage of cuda_version. #9655

Merged

auto-merge was automatically disabled October 2, 2025 14:26
Head branch was pushed to by a user without write access

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from a27f105 to 63bb881 Compare October 2, 2025 14:26

sshonTT force-pushed the sshon/custom-compiler-options-upstream branch from 63bb881 to 9799ed0 Compare October 2, 2025 18:23

zhanyong-wan approved these changes Oct 3, 2025

View reviewed changes

zhanyong-wan enabled auto-merge (squash) October 9, 2025 15:07

zhanyong-wan merged commit b1131d1 into pytorch:master Oct 9, 2025
24 checks passed

sshonTT deleted the sshon/custom-compiler-options-upstream branch October 9, 2025 15:53

feat: add support for custom compile options in torch_xla.compile and… #9575

feat: add support for custom compile options in torch_xla.compile and… #9575

Conversation

sshonTT commented Aug 21, 2025

Uh oh!

sshonTT commented Aug 21, 2025

Uh oh!

qihqi commented Aug 25, 2025

Uh oh!

sshonTT commented Aug 26, 2025

Uh oh!

sshonTT commented Aug 29, 2025

Uh oh!

sshonTT commented Sep 2, 2025

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sshonTT commented Sep 11, 2025

Uh oh!

sshonTT commented Oct 2, 2025

Uh oh!

sshonTT commented Oct 2, 2025

Uh oh!

sshonTT commented Oct 7, 2025

Uh oh!

sshonTT commented Oct 9, 2025

Uh oh!

zhanyong-wan commented Oct 9, 2025

Uh oh!

Uh oh!

Uh oh!