Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[IFRT] add c-bindings for "Held" PjRt classes #751

Merged
merged 9 commits into from
Feb 19, 2025

Conversation

mofeing
Copy link
Collaborator

@mofeing mofeing commented Feb 15, 2025

@wsmoses this PR should solve the issues that arised in #738

@avik-pal i've added ifrt_client_make_single_shard_array_from_host_buffer for creating single-shard arrays and ifrt_client_assemble_array_from_single_shards for joining them into a sharded array. ifrt_pjrt_ArrayFromHostBuffer (now renamed to ifrt_pjrt_array_create) is "IFRT-PjRt" only and wouldn't work for IFRT-proxy.

the only missing thing to use ifrt_client_make_single_shard_array_from_host_buffer is that we need to pass a ifrt::Device*...

Comment on lines 1566 to 1574
auto options = std::make_unique<CompileOptions>(
GenerateCompileOptions(
device_id,
is_sharded,
mesh_ids,
num_mesh_ids,
xla_gpu_cuda_data_dir
)
);
Copy link
Collaborator Author

@mofeing mofeing Feb 15, 2025

Choose a reason for hiding this comment

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

this is crashing because compiler->Compile is expecting a ifrt::Compiler::CompileOptions and not xla::CompileOptions

@avik-pal would you mind taking care of this? i think we need to use this https://github.com/openxla/xla/blob/81335eabdc55f0a2ff02bde7f79da7adda7af2c9/xla/python/ifrt/ir/ifrt_ir_program.h#L102-L109

@mofeing
Copy link
Collaborator Author

mofeing commented Feb 17, 2025

Do not merge. See #751 (comment)

@mofeing mofeing marked this pull request as draft February 17, 2025 09:24
@mofeing mofeing mentioned this pull request Feb 18, 2025
5 tasks
@mofeing mofeing marked this pull request as ready for review February 19, 2025 11:40
@mofeing mofeing merged commit acaabc4 into main Feb 19, 2025
30 of 33 checks passed
@mofeing mofeing deleted the ss/ifrt-held-buffer-c-api branch February 19, 2025 16:25
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