-
Notifications
You must be signed in to change notification settings - Fork 20
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
Conversation
deps/ReactantExtra/API.cpp
Outdated
auto options = std::make_unique<CompileOptions>( | ||
GenerateCompileOptions( | ||
device_id, | ||
is_sharded, | ||
mesh_ids, | ||
num_mesh_ids, | ||
xla_gpu_cuda_data_dir | ||
) | ||
); |
There was a problem hiding this comment.
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
Do not merge. See #751 (comment) |
@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 andifrt_client_assemble_array_from_single_shards
for joining them into a sharded array.ifrt_pjrt_ArrayFromHostBuffer
(now renamed toifrt_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 aifrt::Device*
...