#sdy define sdy::NamedComputationOp.

[copybara-service]

#sdy define sdy::NamedComputationOp.

In JAX host-offloading, XLA adds/uses information to backend_config throughout the compiler, and XLA doesn't inline kCalls with backend_configs*. And WRT to GSPMD, GSPMD is able to propagate through to the function itself. For example, here is what GSPMD propagation returns for a host offloading TPU program:

HloModule jit_f, entry_computation_layout={(s32[8,2]{0,1:T(2,128)S(5)})->s32[8,2]{0,1:T(2,128)S(5)}}, num_partitions=4

%g.2 (Arg_0.3: s32[8,2]) -> s32[8,2] {
  %Arg_0.3 = s32[8,2]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="jit(f)/jit(main)/pjit"}
  ROOT %multiply.4 = s32[8,2]{1,0} multiply(s32[8,2]{1,0} %Arg_0.3, s32[8,2]{1,0} %Arg_0.3), sharding={devices=[2,2]<=[4]}, frontend_attributes={_xla_compute_type="host"}, metadata={op_name="jit(f)/jit(main)/jit(g)/mul" source_file="third_party/py/jax/tests/memories_test.py" source_line=1017}
}

ENTRY %main.9 (Arg_0.1: s32[8,2]) -> s32[8,2] {
  %Arg_0.1 = s32[8,2]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="x"}
  %call.5 = s32[8,2]{1,0} call(s32[8,2]{1,0} %Arg_0.1), to_apply=%g.2, sharding={devices=[2,2]<=[4]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"device_type":"DEVICE_TYPE_HOST","used_scoped_memory_configs":[]}
  ROOT %custom-call = s32[8,2]{1,0} custom-call(s32[8,2]{1,0} %call.5), custom_call_target="MoveToHost", sharding={devices=[2,2]<=[4]}
}

So to handle this, we will introduce a simple region based op that holds the body of the function being called. We will also attach a name to it which will be used during XLA export to give it the original function name.