Lower _linalg_svd

test/cpp/test_aten_xla_tensor.cpp

@@ -913,6 +913,38 @@ TEST_F(AtenXlaTensorTest, TestSVD) {
  });
  }
  }
+ ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
+ ExpectCounterChanged("xla::svd", cpp_test::GetIgnoredCounters());
+}
+
+TEST_F(AtenXlaTensorTest, TestLinalgSVD) {
+ static const int dims[] = {4, 7};
+ for (auto m : dims) {
+ for (auto n : dims) {
+ torch::Tensor a =
+ torch::rand({m, n}, torch::TensorOptions(torch::kFloat));
+ auto b =
+ torch::_linalg_svd(a, /*full_matrices=*/false, /*compute_uv=*/true);
+ ForEachDevice([&](const torch::Device& device) {
+ torch::Tensor xla_a = CopyToDevice(a, device);
+ auto xla_b = torch::_linalg_svd(xla_a, /*full_matrices=*/false,
+ /*compute_uv=*/true);
+ // The U and V matrices might have different sign for column vectors, so
+ // cannot be compared if not by absolute value.
+ AllClose(std::get<0>(b).abs(), std::get<0>(xla_b).abs(), /*rtol=*/1e-3,
+ /*atol=*/1e-4);
+ torch::Tensor diag = std::get<1>(b);
+ torch::Tensor xla_diag = std::get<1>(xla_b);
+ ASSERT_EQ(diag.sizes(), xla_diag.sizes());
+ AllClose(diag, xla_diag, /*rtol=*/1e-3,
+ /*atol=*/1e-4);
+ AllClose(std::get<2>(b).abs(), std::get<2>(xla_b).abs(), /*rtol=*/1e-3,
+ /*atol=*/1e-4);
+ });
+ }
+ }
+ ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
+ ExpectCounterChanged("xla::_linalg_svd", cpp_test::GetIgnoredCounters());
 }
 
 TEST_F(AtenXlaTensorTest, TestQR) {

test/test_ops.py

@@ -48,7 +48,6 @@ def __new__(cls, name, variant_test_name=""):
  AllowedOpInfoEntry('bitwise_right_shift'),
  AllowedOpInfoEntry('ceil'),
  AllowedOpInfoEntry('cholesky'),
- AllowedOpInfoEntry('cholesky_inverse'),
  AllowedOpInfoEntry('chunk'),
  AllowedOpInfoEntry('clone'),
  AllowedOpInfoEntry('contiguous'),
@@ -100,20 +99,14 @@ def __new__(cls, name, variant_test_name=""):
  AllowedOpInfoEntry('floor_divide'),
  AllowedOpInfoEntry('frexp'),
  AllowedOpInfoEntry('ge'),
- AllowedOpInfoEntry('geqrf'),
  AllowedOpInfoEntry('gt'),
  AllowedOpInfoEntry('imag'),
  AllowedOpInfoEntry('inverse'),
  AllowedOpInfoEntry('isin'),
  AllowedOpInfoEntry('le'),
- AllowedOpInfoEntry('linalg.det'),
  AllowedOpInfoEntry('linalg.cholesky'),
  AllowedOpInfoEntry('linalg.cholesky_ex'),
- AllowedOpInfoEntry('linalg.eig'),
  AllowedOpInfoEntry('linalg.householder_product'),
- AllowedOpInfoEntry('linalg.matrix_power'),
- AllowedOpInfoEntry('linalg.qr'),
- AllowedOpInfoEntry('linalg.slogdet'),
  AllowedOpInfoEntry('log'),
  AllowedOpInfoEntry('log10'),
  AllowedOpInfoEntry('log1p'),
@@ -122,7 +115,6 @@ def __new__(cls, name, variant_test_name=""):
  AllowedOpInfoEntry('logaddexp2'),
  AllowedOpInfoEntry('logical_not'),
  AllowedOpInfoEntry('lt'),
- AllowedOpInfoEntry('lu'),
  AllowedOpInfoEntry('lu_unpack'),
  AllowedOpInfoEntry('masked_fill'),
  AllowedOpInfoEntry('masked_scatter'),
@@ -157,7 +149,6 @@ def __new__(cls, name, variant_test_name=""):
  AllowedOpInfoEntry('permute'),
  AllowedOpInfoEntry('pow'),
  AllowedOpInfoEntry('float_power'),
- AllowedOpInfoEntry('qr'),
  AllowedOpInfoEntry('rad2deg'),
  AllowedOpInfoEntry('real'),
  AllowedOpInfoEntry('roll'),
@@ -181,19 +172,12 @@ def __new__(cls, name, variant_test_name=""):
  AllowedOpInfoEntry('hsplit'),
  AllowedOpInfoEntry('vsplit'),
  AllowedOpInfoEntry('dsplit'),
- AllowedOpInfoEntry('triangular_solve'),
  AllowedOpInfoEntry('trunc'),
  AllowedOpInfoEntry('exp2'),
  AllowedOpInfoEntry('nan_to_num'),
  AllowedOpInfoEntry('square'),
  AllowedOpInfoEntry('lerp'),
- AllowedOpInfoEntry('linalg.inv'),
- AllowedOpInfoEntry('linalg.inv_ex'),
  AllowedOpInfoEntry('angle'),
- AllowedOpInfoEntry('linalg.solve'),
- AllowedOpInfoEntry('linalg.matrix_rank'),
- AllowedOpInfoEntry('linalg.svd'),
- AllowedOpInfoEntry('linalg.svdvals'),
  AllowedOpInfoEntry('polar'),
  AllowedOpInfoEntry('ravel'),
  AllowedOpInfoEntry('reshape'),
@@ -276,6 +260,7 @@ def __new__(cls, name, variant_test_name=""):
  # AllowedOpInfoEntry('asinh'),
  # AllowedOpInfoEntry('atan'),
  # AllowedOpInfoEntry('atanh'),
+ # AllowedOpInfoEntry('cholesky_inverse'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('cos'),
  # AllowedOpInfoEntry('cosh'),
  # AllowedOpInfoEntry('cov'),
@@ -287,15 +272,28 @@ def __new__(cls, name, variant_test_name=""):
  # AllowedOpInfoEntry('diagonal'),
  # AllowedOpInfoEntry('fft.ifft'),
  # AllowedOpInfoEntry('fft.ihfft'),
+ # AllowedOpInfoEntry('geqrf'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('gradient'),
  # AllowedOpInfoEntry('kthvalue'),
  # AllowedOpInfoEntry('linalg.cond'),
+ # AllowedOpInfoEntry('linalg.det'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.eig'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('linalg.eigh'),
  # AllowedOpInfoEntry('linalg.eigvalsh'),
+ # AllowedOpInfoEntry('linalg.inv'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.inv_ex'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.slogdet'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.qr'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('linalg.lstsq'),
  # AllowedOpInfoEntry('linalg.norm'),
  # AllowedOpInfoEntry('linalg.matrix_norm'),
+ # AllowedOpInfoEntry('linalg.matrix_rank'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.matrix_power'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.solve'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.svd'), # Slice dim size 1 greater than dynamic slice dimension: 0
+ # AllowedOpInfoEntry('linalg.svdvals'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('linalg.vector_norm'),
+ # AllowedOpInfoEntry('lu'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('std_mean'),
  # AllowedOpInfoEntry('sum'),
  # AllowedOpInfoEntry('mean'),
@@ -328,6 +326,7 @@ def __new__(cls, name, variant_test_name=""):
  # AllowedOpInfoEntry('repeat'),
  # AllowedOpInfoEntry('squeeze'),
  # AllowedOpInfoEntry('tile'),
+ # AllowedOpInfoEntry('triangular_solve'), # Slice dim size 1 greater than dynamic slice dimension: 0
  # AllowedOpInfoEntry('var'),
  # AllowedOpInfoEntry('logsumexp'),
  # AllowedOpInfoEntry('transpose'),
@@ -339,6 +338,7 @@ def __new__(cls, name, variant_test_name=""):
  # AllowedOpInfoEntry('norm'),
  # AllowedOpInfoEntry('t'),
  # AllowedOpInfoEntry('logdet'), xla::lodget does not handle empty input
+ # AllowedOpInfoEntry('qr'), # Slice dim size 1 greater than dynamic slice dimension: 0
 
  # Failed on CUDA CI only (investigate)
  # app.circleci.com/pipelines/github/pytorch/xla/9088/workflows/2d59c649-db2b-4384-921e-5e43eba1b51a/jobs/17875

torch_xla/csrc/aten_xla_type.cpp

@@ -3032,6 +3032,32 @@ at::Tensor& XLANativeFunctions::zero_(at::Tensor& self) {
  return self;
 }
 
+std::tuple<at::Tensor, at::Tensor, at::Tensor> XLANativeFunctions::_linalg_svd(
+ const at::Tensor& self, bool full_matrices, bool compute_uv,
+ c10::optional<c10::string_view> /* driver */) {
+ // The optional driver string is only for CUDA with a cuSOLVER backend.
+ TORCH_LAZY_FN_COUNTER("xla::");
+ // As per https://pytorch.org/docs/stable/generated/torch.svd.html,
+ // The second boolean argument is exactly opposite between
+ // torch::svd and torch::_linalg_svd, hence the negation of full_matrices.
+ XLATensorPtr self_tensor = bridge::GetXlaTensor(self);
+ auto results = tensor_methods::svd(self_tensor, !full_matrices, compute_uv);
+ auto u = std::get<0>(results);
+ auto s = std::get<1>(results);
+ auto vh = tensor_methods::transpose(std::get<2>(results), 0, 1);
+ if (!compute_uv) {
+ // When compute_uv is false, torch::_linalg_svd returns an empty tensor for
+ // u and vh.
+ u = tensor_methods::full({0}, 0, self_tensor->GetDevice(),
+ self_tensor->dtype());
+ vh = tensor_methods::full({0}, 0, self_tensor->GetDevice(),
+ self_tensor->dtype());
+ }
+ return std::make_tuple(bridge::AtenFromXlaTensor(u),
+ bridge::AtenFromXlaTensor(s),
+ bridge::AtenFromXlaTensor(vh));
+}
+
 at::Scalar XLANativeFunctions::_local_scalar_dense(const at::Tensor& self) {
  if (DebugUtil::ExperimentEnabled("early_sync")) {
  // sync tensors in order to save computation when step is marked later.

xla_native_functions.yaml

@@ -112,6 +112,7 @@ supported:
  - _copy_from
  - _copy_from_and_resize
  - _index_put_impl_
+ - _linalg_svd
  - _local_scalar_dense
  - _log_softmax
  - _log_softmax_backward_data