[XLA:TPU] Support output streaming and refactor TryOutputStreaming in…

…to a bottoms-up approach.

Previously, output streaming took a top-down approach which indiscriminately checks if a MoveToHost custom call would trace down to an output marked with host memory space. This did not work when a dynamic-update-slice existed between the MTH call and the output. This CL fixes this problem by handling output streaming before other MTH calls, while also improving efficiency with the bottoms-up approach so we only trace a single path in the graph.

PiperOrigin-RevId: 630885979

tests/memories_test.py

@@ -1195,6 +1195,22 @@ def body(carry, x):
     out_s = NamedSharding(mesh, P(None, None, "z"), memory_kind="device")
     self.assertEqual(out_hbm.sharding, out_s)
 
+  def test_output_streaming(self):
+    mesh = jtu.create_global_mesh((1, 1), ("x", "y"))
+    np_inp = np.arange(16.0).reshape(8, 2)
+    s_hbm = NamedSharding(mesh, P("x", "y"), memory_kind="device")
+    s_host = NamedSharding(mesh, P("x", "y"), memory_kind="pinned_host")
+    arr_hbm = jax.device_put(np_inp, s_hbm)
+
+    @functools.partial(jax.jit, out_shardings=s_host)
+    def f(xs):
+      out_tpu = xs + 1.0
+      return out_tpu
+
+    out_host = f(arr_hbm)
+    self.assertArraysEqual(out_host, np_inp + 1.0)
+    self.assertEqual(out_host.sharding, s_host)
+
 
 class ActivationOffloadingTest(jtu.JaxTestCase):