Introduce HLO graph bindings

[rpsilva-aws]

In this PR, we introduce the means to print the protobuf of an HLO graph. This is helpful to debug hash-related issues with the caching that depends on the serialized proto. It will also add the flexibility to use the same hash to manipulate the hash (e.g. with pytorch/pytorch#144489). Once the flexibility to do so is merged on PyTorch, we can introduce this in our tests to test various applications of the cache in a single in-memory execution.

Example:

    device = xm.xla_device()
    a = torch.zeros(2048, device=device, requires_grad=True)
    xs.mark_sharding(a, spmd_mesh, ('x',))
    b = torch.randn([32, 2048], device=device, requires_grad=True)
    xs.mark_sharding(b, spmd_mesh, (None, 'y'))

    device = xm.xla_device()
    a = torch.zeros(2048, device=device, requires_grad=True)
    xs.mark_sharding(a, spmd_mesh, ('x',))
    b = torch.randn([32, 2048], device=device, requires_grad=True)
    xs.mark_sharding(b, spmd_mesh, (None, 'y'))

    def fn(x, y):
      x = x + 1
      return x, y * 2

    result = fn(a, b)
    hlo_graph_bytecode = torch_xla._XLAC._get_xla_tensors_hlo_bytecode(result)

b'\n\nIrToHlo.75\x12\nIrToHlo.75\x1a\xa7#\n\nIrToHlo.75\x128\n\nconstant.1\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x80?\x98\x02\x01\xa2\x04\x00\x128\n\nconstant.2\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x80?\x98\x02\x02\xa2\x04\x00\x128\n\nconstant.3\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x00\x00\x98\x02\x03\xa2\x04\x00\x122\n\treshape.4\x12\x07reshape\x1a\x10\x10\x0b\x1a\x01\x01*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02\x04\xa2\x02\x01\x03\xa2\x04\x00\x129\n\x0bbroadcast.5\x12\tbroadcast\x1a\x10\x10\x0b\x1a\x01\x01*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00r\x01\x00\x98\x02\x05\xa2\x02\x01\x04\xa2\x04\x00\x12)\n\treshape.6\x12\x07reshape\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02\x06\xa2\x02\x01\x05\xa2\x04\x00\x127\n\x0bbroadcast.7\x12\tbroadcast\x1a\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02\x07\xa2\x02\x01\x06\xa2\x04\x00\x12L\n\rcustom-call.8\x12\x0bcustom-call\x1a\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\xe2\x01\x08Sharding\x98\x02\x08\xa2\x02\x01\x07\xc2\x02\x00\xa2\x04\x00\xe8\x04\x01\x12,\n\nmultiply.9\x12\x08multiply\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02\t\xa2\x02\x02\x02\x01\xa2\x04\x00\x128\n\x0cbroadcast.10\x12\tbroadcast\x1a\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02\n\xa2\x02\x01\t\xa2\x04\x00\x12-\n\x06add.11\x12\x03add\x1a\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02\x0b\xa2\x02\x02\x08\n\xa2\x04\x00\x12&\n\x05p0.12\x12\tparameter\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02\x0c\xc2\x02\x00\xa2\x04\x00\x12(\n\x05p1.13\x12\tparameter\x1a\x07\x10\x05*\x03\x80\x01\x01:\x00H\x01\x98\x02\r\xc2\x02\x00\xa2\x04\x00\x128\n\x0bconstant.14\x12\x08constant\x1a\x07\x10\x05*\x03\x80\x01\x01:\x00B\x0e\n\x07\x10\x05*\x03\x80\x01\x01*\x03\xfd\x87\r\x98\x02\x0e\xa2\x04\x00\x12-\n\x0bmultiply.15\x12\x08multiply\x1a\x07\x10\x05*\x03\x80\x01\x01:\x00\x98\x02\x0f\xa2\x02\x02\x0e\r\xa2\x04\x00\x129\n\x0bconstant.16\x12\x08constant\x1a\x07\x10\x05*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x05*\x03\x80\x01\x01*\x04\xc3\xbd\x9a\x01\x98\x02\x10\xa2\x04\x00\x12#\n\x06add.17\x12\x03add\x1a\x07\x10\x05*\x03\x80\x01\x01:\x00\x98\x02\x11\xa2\x02\x02\x10\x0f\xa2\x04\x00\x129\n\x0bconstant.18\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x80?\x98\x02\x12\xa2\x04\x00\x126\n\nreshape.19\x12\x07reshape\x1a\x13\x10\x0b\x1a\x02\x01\x01*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02\x13\xa2\x02\x01\x12\xa2\x04\x00\x12>\n\x0cbroadcast.20\x12\tbroadcast\x1a\x13\x10\x0b\x1a\x02\x01\x01*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00r\x02\x00\x01\x98\x02\x14\xa2\x02\x01\x13\xa2\x04\x00\x12*\n\nreshape.21\x12\x07reshape\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02\x15\xa2\x02\x01\x14\xa2\x04\x00\x12;\n\x0cbroadcast.22\x12\tbroadcast\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02\x16\xa2\x02\x01\x15\xa2\x04\x00\x129\n\x0bconstant.23\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x00\x00\x98\x02\x17\xa2\x04\x00\x126\n\nreshape.24\x12\x07reshape\x1a\x13\x10\x0b\x1a\x02\x01\x01*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02\x18\xa2\x02\x01\x17\xa2\x04\x00\x12>\n\x0cbroadcast.25\x12\tbroadcast\x1a\x13\x10\x0b\x1a\x02\x01\x01*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00r\x02\x00\x01\x98\x02\x19\xa2\x02\x01\x18\xa2\x04\x00\x12*\n\nreshape.26\x12\x07reshape\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02\x1a\xa2\x02\x01\x19\xa2\x04\x00\x12;\n\x0cbroadcast.27\x12\tbroadcast\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02\x1b\xa2\x02\x01\x1a\xa2\x04\x00\x12*\n\nconvert.28\x12\x07convert\x1a\x07\x10\t*\x03\x80\x01\x01:\x00\x98\x02\x1c\xa2\x02\x01\x11\xa2\x04\x00\x126\n\x0bconstant.29\x12\x08constant\x1a\x07\x10\t*\x03\x80\x01\x01:\x00B\x0c\n\x07\x10\t*\x03\x80\x01\x01:\x01\x00\x98\x02\x1d\xa2\x04\x00\x129\n\x0bconstant.30\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x80?\x98\x02\x1e\xa2\x04\x00\x129\n\x0bconstant.31\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x00\x00\x98\x02\x1f\xa2\x04\x00\x123\n\nreshape.32\x12\x07reshape\x1a\x10\x10\t\x1a\x01\x01*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02 \xa2\x02\x01\x1c\xa2\x04\x00\x123\n\nreshape.33\x12\x07reshape\x1a\x10\x10\t\x1a\x01\x01*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02!\xa2\x02\x01\x1d\xa2\x04\x00\x12?\n\x0econcatenate.34\x12\x0bconcatenate\x1a\x10\x10\t\x1a\x01\x02*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00r\x01\x00\x98\x02"\xa2\x02\x02 !\xa2\x04\x00\x12d\n\x14rng-bit-generator.35\x12\x11rng-bit-generator\x1a-\x10\r"\x10\x10\t\x1a\x01\x02*\x06\n\x01\x00\x80\x01\x012\x01\x00"\x17\x10\x08\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02#\xa2\x02\x01"\xa2\x04\x00\x12P\n\x14get-tuple-element.36\x12\x11get-tuple-element\x1a\x17\x10\x08\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00h\x01\x98\x02$\xa2\x02\x01#\xa2\x04\x00\x12G\n\x14get-tuple-element.37\x12\x11get-tuple-element\x1a\x10\x10\t\x1a\x01\x02*\x06\n\x01\x00\x80\x01\x012\x01\x00:\x00\x98\x02%\xa2\x02\x01#\xa2\x04\x00\x126\n\x0bconstant.38\x12\x08constant\x1a\x07\x10\x08*\x03\x80\x01\x01:\x00B\x0c\n\x07\x10\x08*\x03\x80\x01\x012\x01\t\x98\x02&\xa2\x04\x00\x12>\n\x0cbroadcast.39\x12\tbroadcast\x1a\x17\x10\x08\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02\'\xa2\x02\x01&\xa2\x04\x00\x12S\n\x16shift-right-logical.40\x12\x13shift-right-logical\x1a\x17\x10\x08\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02(\xa2\x02\x02$\'\xa2\x04\x00\x12:\n\nconvert.41\x12\x07convert\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02)\xa2\x02\x01(\xa2\x04\x00\x129\n\x0bconstant.42\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x004\x98\x02*\xa2\x04\x00\x12>\n\x0cbroadcast.43\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02+\xa2\x02\x01*\xa2\x04\x00\x12=\n\x0bmultiply.44\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02,\xa2\x02\x02)+\xa2\x04\x00\x12-\n\x0bsubtract.45\x12\x08subtract\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00\x98\x02-\xa2\x02\x02\x1e\x1f\xa2\x04\x00\x12>\n\x0cbroadcast.46\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02.\xa2\x02\x01-\xa2\x04\x00\x12=\n\x0bmultiply.47\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02/\xa2\x02\x02,.\xa2\x04\x00\x12>\n\x0cbroadcast.48\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x020\xa2\x02\x01\x1f\xa2\x04\x00\x123\n\x06add.49\x12\x03add\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x021\xa2\x02\x02/0\xa2\x04\x00\x12L\n\x08slice.50\x12\x05slice\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x8a\x01\x04\x10\x10\x18\x01\x8a\x01\x04\x10\x01\x18\x01\x8a\x01\x05\x10\x80\x10\x18\x01\x98\x022\xa2\x02\x011\xa2\x04\x00\x12N\n\x08slice.51\x12\x05slice\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x8a\x01\x04\x10\x10\x18\x01\x8a\x01\x06\x08\x01\x10\x02\x18\x01\x8a\x01\x05\x10\x80\x10\x18\x01\x98\x023\xa2\x02\x011\xa2\x04\x00\x129\n\x0bconstant.52\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x95\xbf\xd63\x98\x024\xa2\x04\x00\x12>\n\x0cbroadcast.53\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x025\xa2\x02\x014\xa2\x04\x00\x12;\n\nmaximum.54\x12\x07maximum\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x026\xa2\x02\x0225\xa2\x04\x00\x129\n\x0bconstant.55\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\xdb\x0f\xc9@\x98\x027\xa2\x04\x00\x12>\n\x0cbroadcast.56\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x028\xa2\x02\x017\xa2\x04\x00\x12=\n\x0bmultiply.57\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x029\xa2\x02\x0283\xa2\x04\x00\x122\n\x06log.58\x12\x03log\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02:\xa2\x02\x016\xa2\x04\x00\x129\n\x0bconstant.59\x12\x08constant\x1a\x07\x10\x0b*\x03\x80\x01\x01:\x00B\x0f\n\x07\x10\x0b*\x03\x80\x01\x01B\x04\x00\x00\x00\xc0\x98\x02;\xa2\x04\x00\x12>\n\x0cbroadcast.60\x12\tbroadcast\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02<\xa2\x02\x01;\xa2\x04\x00\x12=\n\x0bmultiply.61\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02=\xa2\x02\x02<:\xa2\x04\x00\x124\n\x07sqrt.62\x12\x04sqrt\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02>\xa2\x02\x01=\xa2\x04\x00\x124\n\x07sine.63\x12\x04sine\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02?\xa2\x02\x019\xa2\x04\x00\x12=\n\x0bmultiply.64\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02@\xa2\x02\x02?>\xa2\x04\x00\x128\n\tcosine.65\x12\x06cosine\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02A\xa2\x02\x019\xa2\x04\x00\x12=\n\x0bmultiply.66\x12\x08multiply\x1a\x17\x10\x0b\x1a\x04\x10\x01\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00\x98\x02B\xa2\x02\x02A>\xa2\x04\x00\x12F\n\x0econcatenate.67\x12\x0bconcatenate\x1a\x17\x10\x0b\x1a\x04\x10\x02\x80\x10*\x08\n\x03\x02\x01\x00\x80\x01\x012\x03\x00\x00\x00:\x00r\x01\x01\x98\x02C\xa2\x02\x02@B\xa2\x04\x00\x127\n\nreshape.68\x12\x07reshape\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02D\xa2\x02\x01C\xa2\x04\x00\x12:\n\x0bmultiply.69\x12\x08multiply\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02E\xa2\x02\x02D\x16\xa2\x04\x00\x120\n\x06add.70\x12\x03add\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02F\xa2\x02\x02\x1bE\xa2\x04\x00\x12P\n\x0ecustom-call.71\x12\x0bcustom-call\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\xe2\x01\x08Sharding\x98\x02G\xa2\x02\x01F\xc2\x02\x00\xa2\x04\x00\xe8\x04\x01\x12;\n\x0cbroadcast.72\x12\tbroadcast\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02H\xa2\x02\x01\x0c\xa2\x04\x00\x12:\n\x0bmultiply.73\x12\x08multiply\x1a\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02I\xa2\x02\x02GH\xa2\x04\x00\x12K\n\x08tuple.74\x12\x05tuple\x1a+\x10\r"\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00"\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00:\x00\x98\x02J\xa2\x02\x02\x0bI\xa2\x04\x00"G\n\x07\x10\x0b*\x03\x80\x01\x01\n\x07\x10\x05*\x03\x80\x01\x01\x12+\x10\r"\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00"\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00\x1a\x02p0\x1a\x02p1(K0J"G\n\x07\x10\x0b*\x03\x80\x01\x01\n\x07\x10\x05*\x03\x80\x01\x01\x12+\x10\r"\x11\x10\x0b\x1a\x02\x80\x10*\x06\n\x01\x00\x80\x01\x012\x01\x00"\x14\x10\x0b\x1a\x03 \x80\x10*\x07\n\x02\x01\x00\x80\x01\x012\x02\x00\x00\x1a\x02p0\x1a\x02p1(K0K\x8a\x01\x00'

[rpsilva-aws]

@tengyifei