Split and Merge Hlo

[merrymercy]

Tips

• Call python code in c++
  Example: https://github.com/parax-project/tensorflow-parax/blob/b13493fd4429e9f6f4373c5d2e8503a3ad6d020f/tensorflow/compiler/xla/service/gpu/auto_sharding.cc#L1382-L1417
  Pybind11 doc: https://pybind11.readthedocs.io/en/stable/advanced/embedding.html#

• Where to get the hlo module? Right after the auto-sharding pass
  https://github.com/parax-project/tensorflow-parax/blob/b13493fd4429e9f6f4373c5d2e8503a3ad6d020f/tensorflow/compiler/xla/service/gpu/auto_sharding.cc#L1646-L1648

• Export c++ funtion to python
  Add your function here https://github.com/parax-project/tensorflow-parax/blob/b13493fd4429e9f6f4373c5d2e8503a3ad6d020f/tensorflow/compiler/xla/python/xla_compiler.cc#L178

• A possible implementation
  In auto_sharding.cc (the last line)

  // insert before the statement "return true;"
  forward_module, backward_module = partition_hlo_module(module);

  PyGILState_STATE gstate = PyGILState_Ensure();
  {
    py::object submodule = py::module_::import("parax.auto_sharding");
    py::object set_forward_backward_module = submodule.attr("set_forward_backward_module");
    set_forward_backward_module(forward_module, backward_module);
  }
  PyGILState_Release(gstate);

in parax/auto_sharding.py

received_forward_module = None
received_backward_module = None
def set_forward_backward_module(forward_module, backward_module):
    global received_forward_module, received_backward_module
    received_forward_module = forward_module
    received_backward_module = backward_module

def split_and_compile(forward_hlo, backward_hlo):
    global received_forward_module, received_backward_module
    merged_hlo = merge(forward_hlo, backward_hlo)
    compile_with_auto_sharidng(merged_hlo)  # this will set received_forward_module and received_backward_module
    forward_binary = compile_without_auto_sharding(received_forward_module)
    backward_binary = compile_without_auto_sharding(received_backward_module)
    return forward_binary, backward_binary

[merrymercy]

#46