Add alternative dynamo backend

test/dynamo/test_dynamo.py

@@ -10,6 +10,7 @@
 import torch_xla.core.xla_env_vars as xenv
 from torch_xla import runtime as xr
 import torch_xla.debug.profiler as xp
+from torch_xla._dynamo import dynamo_backend2
 import torch.optim as optim
 import torch.nn as nn
 import torch._dynamo as dynamo
@@ -38,31 +39,33 @@ def _is_on_neuron():
 skipOnNeuron = unittest.skipIf(_is_on_neuron(), 'Not supported on NEURON')
 
 
-class DynamoInPlaceTest(unittest.TestCase):
+class DynamoInPlaceTest(parameterized.TestCase):
 
  def inplace_update(self, a):
  a += 1
  return a
 
- def test_inplace_update_correctness(self):
+ @parameterized.parameters(['openxla', dynamo_backend2.dynamo_backend])
+ def test_inplace_update_correctness(self, backend):
  dynamo_inplace = torch.compile(
- self.inplace_update, backend="openxla", fullgraph=True)
+ self.inplace_update, backend=backend, fullgraph=True)
  t = torch.tensor([0, 1, 2], device=xm.xla_device())
  for i in range(10):
  t = dynamo_inplace(t)
  self.assertTrue(torch.all(torch.eq(t.cpu(), torch.tensor([10, 11, 12]))))
 
 
-class DynamRandomOpTest(unittest.TestCase):
+class DynamRandomOpTest(parameterized.TestCase):
 
  def random_op(self, a):
  return torch.randn(5, 5, device=a.device) + a
 
- def test_random_op_different_result_each_run(self):
+ @parameterized.parameters(['openxla', dynamo_backend2.dynamo_backend])
+ def test_random_op_different_result_each_run(self, backend):
  xm.wait_device_ops()
  met.clear_all()
  dynamo_random_op = torch.compile(
- self.random_op, backend="openxla", fullgraph=True)
+ self.random_op, backend=backend, fullgraph=True)
  t = torch.randn(5, 5).to(xm.xla_device())
  dynamo_res_1 = dynamo_random_op(t)
  dynamo_res_2 = dynamo_random_op(t)
@@ -75,7 +78,7 @@ def test_random_op_different_result_each_run(self):
  self.assertFalse(torch.allclose(dynamo_res_2, dynamo_res_3))
 
 
-class DynamoLTCInteractionTest(unittest.TestCase):
+class DynamoLTCInteractionTest(parameterized.TestCase):
 
  def index_copy_inplace(self, cache, update_indices, xk):
  cache.index_copy_(0, update_indices, xk)
@@ -104,21 +107,22 @@ def test_mark_step_after_dynamo(self):
  xm.wait_device_ops()
  self.assertEqual(current_execute_time, met.metric_data('ExecuteTime')[0])
 
- def test_copy_op(self):
+ @parameterized.parameters(['openxla', dynamo_backend2.dynamo_backend])
+ def test_copy_op(self, backend):
 
  def copy_a_to_b(a):
  res = a.cos()
- copy = torch.ops.aten.copy.default(a, res)
+ copy = torch.ops.aten.copy_.default(a, res)
  return copy
 
  device = torch_xla.device()
- compiled_copy = torch.compile(copy_a_to_b, backend="openxla")
+ compiled_copy = torch.compile(copy_a_to_b, backend=backend)
  a = torch.randn(2, 9).to(device)
  res = compiled_copy(a)
  self.assertTrue(torch.allclose(res, a))
 
 
-class DynamoProfilerTest(unittest.TestCase):
+class DynamoProfilerTest(parameterized.TestCase):
 
  def dummy_fn(self, a):
  return torch.sin(a) + a
@@ -253,11 +257,10 @@ def fn_without_input(device):
  res_xla_dynamo = compiled_fn(device)
  self.assertTrue(torch.allclose(res_cpu, res_xla_dynamo.cpu()))
 
- @parameterized.parameters(
- True,
- False,
- )
- def test_simple_model_with_in_place_ops(self, initialize_on_cuda):
+ @parameterized.product(
+ initialize_on_cuda=[True, False],
+ backend=['openxla', dynamo_backend2.dynamo_backend])
+ def test_simple_model_with_in_place_ops(self, initialize_on_cuda, backend):
 
  class TestModel(nn.Module):
 
@@ -286,7 +289,7 @@ def forward(self, index, copy_tensor, input_tensor, op_name):
 
  cpu_model = TestModel()
  device_model = TestModel(device).to(device)
- compiled_model = torch.compile(device_model, backend='openxla')
+ compiled_model = torch.compile(device_model, backend=backend)
 
  input_tensor = torch.ones(3)
  copy_tensor = torch.rand(5, 3)
@@ -306,11 +309,10 @@ def forward(self, index, copy_tensor, input_tensor, op_name):
  op_name=in_place_op)
  self.assertTrue(torch.allclose(res_cpu, res_device_dynamo.cpu()))
 
- @parameterized.parameters(
- True,
- False,
- )
- def test_einsum(self, initialize_on_cuda):
+ @parameterized.product(
+ initialize_on_cuda=[True, False],
+ backend=['openxla', dynamo_backend2.dynamo_backend])
+ def test_einsum(self, initialize_on_cuda, backend):
  # einsum currently does not have meta function to compute the shape hence
  # will fallback to XLA with FakeTensor as input to infer the output shape.
  def einsum_mm(a, b):
@@ -321,7 +323,7 @@ def einsum_mm(a, b):
  b = torch.randn(4, 4, 4, 4).to(device)
  xm.mark_step()
 
- dynamo_einsum_mm = torch.compile(einsum_mm, backend="openxla")
+ dynamo_einsum_mm = torch.compile(einsum_mm, backend=backend)
  res_device_dynamo = dynamo_einsum_mm(a, b)
  res_device_non_dynamo = einsum_mm(a, b)
  self.assertTrue(
@@ -368,11 +370,10 @@ def get_loader(self, device, sample_count, batch_size=4):
 
  @skipOnTpu
  @skipOnNeuron
- @parameterized.parameters(
- True,
- False,
- )
- def test_resnet18(self, initialize_on_cuda):
+ @parameterized.product(
+ initialize_on_cuda=[True, False],
+ backend=['openxla', dynamo_backend2.dynamo_backend])
+ def test_resnet18(self, initialize_on_cuda, backend):
  device = self._choose_proper_device(initialize_on_cuda)
  sample_count = xu.getenv_as('SAMPLE_COUNT', int, defval=10)
  loader = self.get_loader(device, sample_count, batch_size=4)
@@ -386,19 +387,21 @@ def test_resnet18(self, initialize_on_cuda):
  xm.mark_step()
  xm.wait_device_ops()
  met.clear_all()
- dynamo_resnet18 = torch.compile(device_resnet18, backend='openxla')
+ dynamo_resnet18 = torch.compile(device_resnet18, backend=backend)
  for data, _ in loader:
  output = dynamo_resnet18(data)
  output_cpu = resnet18(data.cpu())
  self.assertTrue(
  torch.allclose(output_cpu, output.cpu(), rtol=1e-05, atol=1e-05))
  # We only expect one graph for the resnet18 inference.
- self.assertEqual(met.metric_data('CompileTime')[0], 1)
- self.assertEqual(met.metric_data('ExecuteTime')[0], sample_count)
- self.assertEqual(
- met.metric_data('RunCachedGraphInputData')[0], sample_count)
- self.assertEqual(
- met.metric_data('RunCachedGraphOutputData')[0], sample_count)
+ if backend == 'openxla':
+ # backend2 doesnt populate metrics
+ self.assertEqual(met.metric_data('CompileTime')[0], 1)
+ self.assertEqual(met.metric_data('ExecuteTime')[0], sample_count)
+ self.assertEqual(
+ met.metric_data('RunCachedGraphInputData')[0], sample_count)
+ self.assertEqual(
+ met.metric_data('RunCachedGraphOutputData')[0], sample_count)
 
  @skipOnNeuron
  def test_resnet18_lazy_vs_dynamo(self):
@@ -428,7 +431,7 @@ def test_resnet18_lazy_vs_dynamo(self):
  # mess up the counter check.
 
 
-class DynamoCpuFallbackTest(unittest.TestCase):
+class DynamoCpuFallbackTest(parameterized.TestCase):
 
  def test_operator_fallback(self):
 
@@ -509,7 +512,7 @@ def fn_fallback(t):
  self.assertEqual(met.metric_data('ExecuteTime')[0], 3)
 
 
-class DynamoTrainingBasicTest(unittest.TestCase):
+class DynamoTrainingBasicTest(parameterized.TestCase):
 
  @classmethod
  def setUpClass(self):
@@ -613,7 +616,7 @@ def test_resnet18(self):
  met.metric_data('RunCachedGraphOutputData')[0], sample_count * 2)
 
 
-class DynamoTrainingOptimizerTest(unittest.TestCase):
+class DynamoTrainingOptimizerTest(parameterized.TestCase):
 
  @classmethod
  def setUpClass(self):
@@ -719,7 +722,7 @@ def test_resnet18(self):
  met.metric_data('RunCachedGraphOutputData')[0], sample_count * 3)
 
 
-class DynamoErrorMessageTest(unittest.TestCase):
+class DynamoErrorMessageTest(parameterized.TestCase):
 
  def test_mixed_cpu_tensor(self):
  device = xm.xla_device()
@@ -758,17 +761,18 @@ def test_all_cpu_tensor(self):
  self.assertLessEqual(len(met.counter_names()), 1)
 
 
-class DynamoOperationsTests(test_utils.XlaTestCase):
+class DynamoOperationsTest(test_utils.XlaTestCase, parameterized.TestCase):
 
- def test_new_with_sizes(self):
+ @parameterized.parameters(['openxla', dynamo_backend2.dynamo_backend])
+ def test_new_with_sizes(self, backend):
 
  # The addition operation is needed here, since the error only occurs when FakeTensorMode
  # checks the device of the arguments of some operation. If there's no operation using the
  # result of Tensor.new, this comparison never occurs.
  def foo(x):
  return x.new(*x.size()) + x
 
- optfoo = torch.compile(backend="openxla")(foo)
+ optfoo = torch.compile(backend=backend)(foo)
 
  t = torch.arange(9)
  Xt = t.to(xm.xla_device())
@@ -782,12 +786,13 @@ def foo(x):
  self.assertEqual(expected.dtype, actual.dtype)
  self.assertEqual(expected.device, actual.device)
 
- def test_return_expand(self):
+ @parameterized.parameters(['openxla', dynamo_backend2.dynamo_backend])
+ def test_return_expand(self, backend):
 
  def foo(x):
  return x.expand(2, -1)
 
- optfoo = torch.compile(backend="openxla")(foo)
+ optfoo = torch.compile(backend=backend)(foo)
 
  t = torch.arange(10)
  Xt = t.to(xm.xla_device())

torch_xla/_dynamo/dynamo_backend2.py

@@ -0,0 +1,64 @@
+import functools
+from typing import Any
+import torch
+from torch.utils import _pytree as pytree
+from torch_xla.core import xla_builder as xb
+import torch_xla
+
+from torch._dynamo.backends.common import aot_autograd
+from functorch.compile import make_boxed_func
+
+
+def _dynamo_backend(model: torch.fx.GraphModule, sample_args: Any):
+ """A dynamo backend that compiles a FX graph to HLO using JAX and torchax.
+
+ It takes FX graph as input and returns a compiled PyTorch function. The FX graph
+ is traced into a JAX function using torchax, and the JAX function is lowered to HLO.
+
+ Args:
+ model: the graph to be compiled
+ sample_args: a tuple or list of sample inputs. I.e. model(*sample_args) produces
+ the model output
+
+ Returns:
+ Another callable f such that f(*sample_inputs) computes the same thing as model.
+ """
+
+ try:
+ import torchax.interop
+ from torchax.export import JaxInterpreter
+ import jax
+ except ImportError:
+ print('To use this dynamo backend, please install torchax')
+ raise
+
+ jax.config.update("jax_enable_x64", True)
+ env = torchax.default_env()
+ xla_device = torch_xla.device()
+
+ def run_jax(*args, initial_rng_key):
+ args_t = torchax.interop.torch_view(args)
+ env.manual_seed(initial_rng_key)
+ with env:
+ res = model(*args_t)
+ return torchax.interop.jax_view(res)
+
+ initial_rng_key = torch.tensor(0, device=xla_device, dtype=torch.uint32)
+ computation = xb.jax_func_to_xla_computation(
+ run_jax, sample_args, {'initial_rng_key': initial_rng_key}, 'dynamo_jax')
+
+ def equivalent(*args, **kwargs):
+ kwargs['initial_rng_key'] = torch.randint(
+ 0, 2**32, (), dtype=torch.uint32, device=xla_device)
+ flattened, _ = pytree.tree_flatten((args, kwargs))
+ res = computation(flattened)
+ if not isinstance(res, (list, tuple)):
+ return (res,)
+ return res
+
+ return make_boxed_func(equivalent)
+
+
+def dynamo_backend(fx, args):
+ from functorch.compile import aot_function
+ return aot_function(fx, fw_compiler=_dynamo_backend)

torch_xla/core/xla_builder.py

@@ -911,8 +911,8 @@ def get_hlo():
  import torch_xla.debug.profiler as xp
  # If we see this trace span in the profiler, we'll know that there's a cache miss.
  with xp.Trace('jax_to_hlo'):
- hlo_ir = jax.jit(
-  fn, keep_unused=True).lower(*sample_tensor_args).compiler_ir('hlo')
+ lowered = jax.jit(fn, keep_unused=True).lower(*sample_tensor_args)
+ hlo_ir = lowered.compiler_ir('hlo')
 
  # Get a protobuf representation of the HLO. `as_serialized_hlo_module_proto` is
  # mentioned at https://github.com/jax-ml/jax/discussions/22266

torchax/test/test_context.py

@@ -39,23 +39,23 @@ def test_mode_decorator(self):
 
  def test_same_manual_seed(self):
  with xla_env:
- torch.manual_seed(1234)
+ xla_env.manual_seed(1234)
  x = torch.randn((3, 3))
  self.assertIsInstance(x, tensor.Tensor)
 
- torch.manual_seed(1234)
+ xla_env.manual_seed(1234)
  y = torch.randn((3, 3))
  self.assertIsInstance(y, tensor.Tensor)
 
  self.assertTrue(torch.equal(torchax.tensor.j2t(x._elem), torchax.tensor.j2t(y._elem)))
 
  def test_different_manual_seed(self):
  with xla_env:
- torch.manual_seed(1234)
+ xla_env.manual_seed(1234)
  x = torch.randn((3, 3))
  self.assertIsInstance(x, tensor.Tensor)
 
- torch.manual_seed(12345)
+ xla_env.manual_seed(12345)
  y = torch.randn((3, 3))
  self.assertIsInstance(y, tensor.Tensor)
 

torchax/torchax/ops/jaten.py

@@ -1351,6 +1351,13 @@ def reduce_fn(a, b):
 
  return y, indices
 
+try:
+ @op(torch.ops.xla.max_pool2d_forward)
+ def _xla_max_pool2d_foward(*args, **kwargs):
+ return _aten_max_pool2d_with_indices(*args, **kwargs)[0]
+except AttributeError:
+ pass
+
 
 # TODO add more ops