Use upstream trilinear decomp by default

test/test_operations.py

@@ -2010,6 +2010,19 @@ def foo(x):
  for dtype in test_dtypes:
  test(dtype)
 
+ def test_trilinear_interpolate(self):
+
+ def func(input_volume):
+ output_size = (32, 64, 64)
+ return F.interpolate(
+ input_volume, size=output_size, mode='trilinear', align_corners=False)
+
+ device = torch_xla.device()
+ input_volume = torch.randn(1, 3, 16, 32, 32).to(device)
+ met.clear_all()
+ self.runAtenTest((input_volume), func)
+ assert len(torch_xla._XLAC._get_executed_fallback_ops()) == 0
+
  def test_gelu_backward_different_types(self):
 
  def foo(grad, inp):

torch_xla/__init__.py

@@ -233,5 +233,5 @@ def _init_xla_lazy_backend():
 
 from .torch_xla import *
 
-# register all custom kenels by default
-from .core import custom_kernel
+# register all custom kenels and decomp by default
+from .core import custom_kernel, decomp_registration

torch_xla/core/decomp_registration.py

@@ -0,0 +1,4 @@
+import torch
+
+torch.library.register_kernel("aten::upsample_trilinear3d", "xla",
+ torch._decomp.decompositions.upsample_trilinear3d)