IFRT prototype

[will-cromar]

• Implement new ComputationClient with IFRT, which currently just wraps PJRT.
  • Enable IFRT wrapper with XLA_USE_IFRT=1
• Move PJRT client initialization into initialize_pjrt.cc/h since IFRT wraps the same PjRtClient.
• Minor fixes to existing pjrt_computation_client:
  • Make spmd_device_str const
  • Change the owned PjRtClient to a unique_ptr. Only SE:TPU required us to use shared_ptr, which is now removed.

There's still some opportunities to refactor common functionality up to ComputationClient, but I'm trying to minimize changes to the high-level API for this PR.

IFRT is still highly experimental. Use at your own risk. See my comments below for caveats and limitations.

[will-cromar]

At this point, I am able to run our SPMD ResNet50 example successfully, but it's extremely slow. Any case where I resharded an array, I chose "copy" semantics for safety. I'll have to take another pass and carefully think through the ownership of the underlying data, since the excessive copies are likely contributing to poor performance.

Sharded execution does not appear to be implemented at all within the IFRT/PJRT wrapper, so I marked that method unimplemented for now. Likewise, I xla::ifrt::HloSharding (a wrapper over xla::HloSharding) is not implemented in the wrapper, so we'll have to carry the xla::OpSharding in our own object for now.

Dynamic shape is currently unsupported by xla::ifrt::Shape.

Until we have feature parity, I will keep IFRT as a separate ComputationClient implementation so we don't break PJRT.

[will-cromar]

Coming back to this PR (finally) after merging supporting changes in separate PRs. Performance is significantly better after rebasing -- it only lags PJRT by ~10% on ResNet50 now, compared to 80% in my first draft. There's still room for optimization, particularly around reducing the number of copies used when transforming IFRT arrays.

IFRT is still highly experimental in this state. Known outstanding issues other than performance:

• The IFRT/PJRT wrapper now natively supports xla::HloSharding, but I'm still getting errors at the current XLA pin. I'd like to try again after another pin update since this support was only added in October.
• Sharded execution (aka multiprocess) is completely unsupported in IFRT.
• Dynamic shapes can't be represented by xla::ifrt::Shape.

I'll clean up this PR and send it for review as an optional/experimental setting.

[will-cromar]

Performance on LLama 7B is not bad! It's somewhere between PJRT now and PJRT before I started working on some optimizations this month:

Totally decoded 1007 tokens in 7.03615 seconds

[JackCaoG]

If you don't intend to merge this for 2.2 release, I will hold on the review until the branch cut.

[will-cromar]

Merging after the cut sounds good to me. This won't be useful in the 2.2 release.

[JackCaoG]

I will take a look today

[JackCaoG]

for these utils, can we share between pjrt and ifrt? so they are actually different?

[will-cromar]

This was intentional. I wanted to minimize changes to common code, and ideally we would be able to remove the PJRT computation client at some point. Would you prefer I try to factor out all common functionality in this PR?

[JackCaoG]

let's do a grep of pjrt in this file and replace them with ifrt. Through I am curious did you intend to query the EnvPjrtDevice here?

[will-cromar]

Yeah, see #5677 (comment)

[JackCaoG]

is there a legit use case of shards.zie() == 0?

[will-cromar]

This is how sharded data placeholders get created right now:

xla/torch_xla/csrc/xla_sharding_util.cpp

Lines 590 to 593 in a2f80e4

	auto sharded_data_placeholder =
	runtime::GetComputationClient()->WrapDataShards(
	{}, GetVirtualDevice().toString(), sharding_specs[i]->shape,
	sharding_specs[i]->sharding);

I'd rather update CreateDataPlaceholder to take a sharding and make it more explicit