[PJRT] Expose `ExecutionContext` when executing a `LoadedExecutable` #20429

Corendos · 2024-12-11T15:21:17Z

Currently, there is no way to provide an ExecutionContext when launching an executable. This prevents using it in custom calls through FFI.

I exposed the field in the PJRT_ExecuteOptions struct and wired it in xla::ExecuteOptions

xla/pjrt/c/pjrt_c_api.h

Co-authored-by: Steeve Morin <[email protected]>

xla/pjrt/c/pjrt_c_api.h

penpornk

Could you please help add a unit test for this as well?

ezhulenev · 2024-12-17T21:54:33Z

I'm not sure anything useful can be added as a test, but PR needs a rebase because I bumped pjrt version yesteday

Corendos · 2024-12-18T09:38:29Z

but PR needs a rebase because I bumped pjrt version yesteday

Addressed in 5b97f17

penpornk

Thank you!

…xecutable` Imported from GitHub PR #20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c21 by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0e by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20429 from Corendos:corendos/pjrt-execution-context 5b97f17 PiperOrigin-RevId: 707074459

…xecutable` Imported from GitHub PR openxla/xla#20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c1ed76cc52dd2b99882cd6dc8d0acc7dbc by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c2146239558ecae141d90aaf707823f89a7d by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0ed720492ecf8d47cafd8d2b09d7222cbf8 by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20429 from Corendos:corendos/pjrt-execution-context 5b97f17549fc5ec0671d19f3fe86a598973f5b61 PiperOrigin-RevId: 707074459

…xecutable` Imported from GitHub PR #20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c21 by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0e by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20429 from Corendos:corendos/pjrt-execution-context 5b97f17 PiperOrigin-RevId: 707074459

…xecutable` Imported from GitHub PR openxla/xla#20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c1ed76cc52dd2b99882cd6dc8d0acc7dbc by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c2146239558ecae141d90aaf707823f89a7d by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0ed720492ecf8d47cafd8d2b09d7222cbf8 by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20429 from Corendos:corendos/pjrt-execution-context 5b97f17549fc5ec0671d19f3fe86a598973f5b61 PiperOrigin-RevId: 707074459

…xecutable` Imported from GitHub PR #20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c21 by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0e by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20429 from Corendos:corendos/pjrt-execution-context 5b97f17 PiperOrigin-RevId: 707074459

…xecutable` Imported from GitHub PR openxla/xla#20429 Currently, there is no way to provide an `ExecutionContext` when launching an executable. This prevents using it in custom calls through FFI. I exposed the field in the `PJRT_ExecuteOptions` struct and wired it in `xla::ExecuteOptions` Copybara import of the project: -- d22969c1ed76cc52dd2b99882cd6dc8d0acc7dbc by Corentin Godeau <[email protected]>: [PJRT] Expose ExecutionContext when executing a LoadedExecutable -- 67d8c2146239558ecae141d90aaf707823f89a7d by Corentin Godeau <[email protected]>: Update xla/pjrt/c/pjrt_c_api.h Co-authored-by: Steeve Morin <[email protected]> -- 5017a0ed720492ecf8d47cafd8d2b09d7222cbf8 by Corentin Godeau <[email protected]>: Added changes to the changelog Merging this change closes #20429 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20429 from Corendos:corendos/pjrt-execution-context 5b97f17549fc5ec0671d19f3fe86a598973f5b61 PiperOrigin-RevId: 707074459

penpornk · 2024-12-18T13:19:53Z

//xla/pjrt:pjrt_c_api_client_test failed. Could you please help take a look? Here is the XLA Linux CPU CI log.

Corendos · 2024-12-19T08:55:07Z

//xla/pjrt:pjrt_c_api_client_test failed. Could you please help take a look? Here is the XLA Linux CPU CI log.

Indeed, thanks ! Addressed in 7128d16

ezhulenev · 2025-01-07T23:23:07Z

I think it's bazel test //folder/..., just run all tests under pjrt, it should be enough

Corendos · 2025-01-08T14:12:30Z

Ok, so it was a bit harder than I thought due to struct visibility.

I had to make pjrt_c_api_client depends on //xla/pjrt/c:pjrt_c_api_wrapper_impl. I'm not sure how it's supposed to be solved properly.

Maybe we should do like the other structures and have a xla::PjRtExecuteContext ?

Corendos · 2025-01-10T16:30:50Z

Ok, I've been thinking about this and I think it all boils down to the fact that, contrary to the other PJRT_* struct, PJRT_ExecuteOptions is not manipulated using a pointer. So we can't just predeclare it and must have access to the struct definition.

One improvement that could be done to avoid being dependent on the _impl.h would be to extract the struct definition elsewhere. Apart from that I don't really know what else could be done. If you have idea @ezhulenev happy to discuss 😁

ezhulenev · 2025-01-10T20:16:42Z

We certainly can't include impl.h headers anywhere :( I started looking at it, and got a bit lost in all our C APIs and structs, will try to come up with something today

ezhulenev · 2025-01-10T21:09:14Z

I think I have a plan, will send PR later today

Context: #20429 PiperOrigin-RevId: 714327866

ezhulenev · 2025-01-11T05:07:34Z

After a few attempts I think this is the one that will work: #21317

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

Context: #20429 PiperOrigin-RevId: 714327866

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

Corendos · 2025-01-13T09:13:36Z

Awesome, I'm looking forward to it being merged !

I guess I'll close this PR then

Context: #20429 PiperOrigin-RevId: 714327866

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: #20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: #20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: #20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: #20429 PiperOrigin-RevId: 715906024

ezhulenev · 2025-01-15T21:00:48Z

Submitted in 7ce927a

+ Correctly (zero/value-)initialize PJRT_ExecuteOptions in tests and pjrt_c_api_client ``` If the number of initializer clauses is less than the number of members or initializer list is completely empty, the remaining members are value-initialized ``` Context: openxla/xla#20429 PiperOrigin-RevId: 715906024

Corendos · 2025-01-16T14:33:47Z

Thanks a lot @ezhulenev !

Corendos and others added 2 commits December 4, 2024 16:14

[PJRT] Expose ExecutionContext when executing a LoadedExecutable

d22969c

Merge branch 'openxla:main' into corendos/pjrt-execution-context

71748c5

penpornk requested a review from ezhulenev December 11, 2024 18:54

ezhulenev approved these changes Dec 12, 2024

View reviewed changes

steeve suggested changes Dec 12, 2024

View reviewed changes

xla/pjrt/c/pjrt_c_api.h Outdated Show resolved Hide resolved

Update xla/pjrt/c/pjrt_c_api.h

67d8c21

Co-authored-by: Steeve Morin <[email protected]>

ezhulenev approved these changes Dec 13, 2024

View reviewed changes

thomasjoerg requested changes Dec 16, 2024

View reviewed changes

xla/pjrt/c/pjrt_c_api.h Outdated Show resolved Hide resolved

Added changes to the changelog

5017a0e

thomasjoerg approved these changes Dec 17, 2024

View reviewed changes

penpornk suggested changes Dec 17, 2024

View reviewed changes

Merge branch 'main' into corendos/pjrt-execution-context

5b97f17

penpornk approved these changes Dec 18, 2024

View reviewed changes

penpornk added the kokoro:force-run Forces CI to rerun label Dec 18, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Dec 18, 2024

copybara-service bot mentioned this pull request Dec 18, 2024

PR #20429: [PJRT] Expose ExecutionContext when executing a LoadedExecutable #20673

Open

copybara-service bot mentioned this pull request Dec 18, 2024

PR #20429: [PJRT] Expose ExecutionContext when executing a LoadedExecutable tensorflow/tensorflow#83253

Draft

Corendos added 2 commits December 19, 2024 09:48

Check that context is non-null before dereferencing it

7128d16

Merge branch 'main' into corendos/pjrt-execution-context

d0b15ef

penpornk approved these changes Dec 19, 2024

View reviewed changes

Corendos force-pushed the corendos/pjrt-execution-context branch from 034e6cc to 4001094 Compare January 8, 2025 14:04

Fix ExecuteOptions marshaling from C++ to C

4e19b2a

Corendos force-pushed the corendos/pjrt-execution-context branch from 4001094 to 4e19b2a Compare January 8, 2025 14:07

copybara-service bot pushed a commit that referenced this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

4ff5821

Context: #20429 PiperOrigin-RevId: 714327866

copybara-service bot mentioned this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client #21317

Merged

copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

4a9244e

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

copybara-service bot mentioned this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client tensorflow/tensorflow#84650

Merged

copybara-service bot pushed a commit that referenced this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

3128032

Context: #20429 PiperOrigin-RevId: 714327866

copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jan 11, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

40172b4

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

copybara-service bot pushed a commit that referenced this pull request Jan 13, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

f75447d

Context: #20429 PiperOrigin-RevId: 714327866

copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jan 13, 2025

[xla:pjrt] Add support for forwarding FFI context to C API client

631c2f7

Context: openxla/xla#20429 PiperOrigin-RevId: 714327866

Corendos closed this Jan 16, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PJRT] Expose `ExecutionContext` when executing a `LoadedExecutable` #20429

[PJRT] Expose `ExecutionContext` when executing a `LoadedExecutable` #20429

Corendos commented Dec 11, 2024

penpornk left a comment

ezhulenev commented Dec 17, 2024

Corendos commented Dec 18, 2024

penpornk left a comment

penpornk commented Dec 18, 2024

Corendos commented Dec 19, 2024

ezhulenev commented Jan 7, 2025

Corendos commented Jan 8, 2025

Corendos commented Jan 10, 2025

ezhulenev commented Jan 10, 2025

ezhulenev commented Jan 10, 2025

ezhulenev commented Jan 11, 2025

Corendos commented Jan 13, 2025

ezhulenev commented Jan 15, 2025

Corendos commented Jan 16, 2025

[PJRT] Expose ExecutionContext when executing a LoadedExecutable #20429

[PJRT] Expose ExecutionContext when executing a LoadedExecutable #20429

Conversation

Corendos commented Dec 11, 2024

penpornk left a comment

Choose a reason for hiding this comment

ezhulenev commented Dec 17, 2024

Corendos commented Dec 18, 2024

penpornk left a comment

Choose a reason for hiding this comment

penpornk commented Dec 18, 2024

Corendos commented Dec 19, 2024

ezhulenev commented Jan 7, 2025

Corendos commented Jan 8, 2025

Corendos commented Jan 10, 2025

ezhulenev commented Jan 10, 2025

ezhulenev commented Jan 10, 2025

ezhulenev commented Jan 11, 2025

Corendos commented Jan 13, 2025

ezhulenev commented Jan 15, 2025

Corendos commented Jan 16, 2025

[PJRT] Expose `ExecutionContext` when executing a `LoadedExecutable` #20429

[PJRT] Expose `ExecutionContext` when executing a `LoadedExecutable` #20429