Load PjRT Plugin API for Custom PjRT Backend

[aws-kingrj]

❓ Questions and Help

We are currently trying to use the load_pjrt_plugin() API for loading custom C-API pjrt plugins but the function is not exposed in torch_xla. I can see it defined here in the xla.cc file of the tensorflow xla commit used for torch_xla 2.0: https://github.com/tensorflow/tensorflow/blob/f7759359f8420d3ca7b9fd19493f2a01bd47b4ef/tensorflow/compiler/xla/python/xla.cc#L325 as a part of the xla_extension.so file as a pybind11 module but that xla_extensions.so file isn't exposed in torch_xla.

Is there another way we can used this API in python to load our C-API pjrt plugin with the current torch_xla 2.0 or is there some other API exposed for us to load pjrt plugins?

[JackCaoG]

@will-cromar Can you take a look?

[will-cromar]

We don't support the xla_extensions pybind module in PyTorch/XLA. You have to load the module in C++ during the ComputationClient creation. For now, you can add a branch like this one here:

xla/third_party/xla_client/pjrt_computation_client.cc

Lines 109 to 113 in bc28ca5

	} else if (device_type == "XPU") {
	TF_VLOG(1) << "Initializing PjRt XPU client...";
	XLA_CHECK_OK(pjrt::LoadPjrtPlugin(
	"xpu", sys_util::GetEnvString(env::kEnvXpuLibraryPath, "libxpu.so")));
	client_ = std::move(xla::GetCApiClient("XPU").value());

You can use the XPU PR as a reference to add a new plugin: #4891

I'm planning on updating this logic to allow dynamic plugin loading and remove all of those device type enums/lists, but that may have to wait until after we remove XRT.

[amithrm]

Thanks @will-cromar ! We will try doing this and update you

[aws-kingrj]

Neuron PR for this #5428

[aws-kingrj]

Neuron PR for import hook to be added: #5429