/* Copyright 2021 The JAX Authors. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include "jaxlib/cpu/lapack_kernels.h" #include #include #include #include #include #include #include #include #include #include "absl/algorithm/container.h" #include "absl/base/dynamic_annotations.h" #include "absl/functional/function_ref.h" #include "absl/status/statusor.h" #include "absl/synchronization/blocking_counter.h" #include "absl/types/span.h" #include "jaxlib/ffi_helpers.h" #include "xla/ffi/api/c_api.h" #include "xla/ffi/api/ffi.h" #include "tsl/platform/platform.h" namespace ffi = xla::ffi; XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::MatrixParams::Side); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::MatrixParams::Transpose); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::MatrixParams::Diag); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::MatrixParams::UpLo); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::svd::ComputationMode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::eig::ComputationMode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::schur::ComputationMode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::schur::Sort); namespace jax { bool lapack_kernels_initialized = false; template inline T CastNoOverflow(int64_t value, std::string_view source = __FILE__) { auto result = MaybeCastNoOverflow(value, source); if (!result.ok()) { throw std::overflow_error{std::string(result.status().message())}; } return result.value(); } template void CopyIfDiffBuffer(ffi::Buffer x, ffi::ResultBuffer x_out) { if (x.typed_data() != x_out->typed_data()) { const auto x_size = x.element_count(); std::copy_n(x.typed_data(), x_size, x_out->typed_data()); } } // Heuristic to choose a good chunk size for parallelizing LAPACK kernels // along the batch dimension. The rough idea is that we want: // * at most as many chunks as we have threads // * we do not want chunks that are so small they would complete in time // comparable to a thread hop. // cost_per_matrix is a FLOP estimate for each batch element. int64_t GetLapackBatchChunkSize(int64_t batch_size, int64_t cost_per_matrix, int64_t num_threads) { #ifdef PLATFORM_GOOGLE // We want the minimum chunk size to be at least as big as the cost of a // thread context switch. Let's assume that's about 2us and each core can do // 100 GFLOP/s. So our minimum chunk size is 2e-6 * 1e11 = 2e5 FLOP. // We could probably tweak this heuristic a bit more. const int64_t kMinWorkPerTask = 200000; if (num_threads <= 1) return batch_size; if (batch_size * cost_per_matrix < kMinWorkPerTask) { return batch_size; } // Size to divide work evenly amongst threads int64_t chunk_for_even_division = (batch_size + num_threads - 1) / num_threads; // Size that ensures each thread gets at least a minimum chunk size. int64_t chunk_for_min_work = (kMinWorkPerTask + cost_per_matrix - 1) / cost_per_matrix; return std::max(chunk_for_even_division, chunk_for_min_work); #else // Parallelizing LAPACK routines over a thread pool is disabled in open-source // builds because it can lead to performance regressions or crashes due to // interactions with external LAPACK/BLAS implementations (e.g., // OpenBLAS/MKL). return batch_size; #endif // PLATFORM_GOOGLE } // Divide batch_count elements into chunk_size pieces, and call run_chunk on // each. Use the thread pool if we have more than one chunk. static ffi::Error ParallelBatchMap( ffi::ThreadPool thread_pool, int64_t batch_count, int64_t chunk_size, absl::FunctionRef run_chunk) { #ifdef PLATFORM_GOOGLE if (chunk_size >= batch_count) { run_chunk(0, batch_count); return ffi::Error::Success(); } else { int64_t num_tasks = (batch_count + chunk_size - 1) / chunk_size; absl::BlockingCounter counter(num_tasks); for (int64_t i = 0; i < batch_count; i += chunk_size) { int64_t current_chunk_size = std::min(chunk_size, batch_count - i); thread_pool.Schedule([run_chunk, &counter, i, current_chunk_size]() { run_chunk(i, current_chunk_size); counter.DecrementCount(); }); } counter.Wait(); return ffi::Error::Success(); } #else // Parallelizing LAPACK routines over a thread pool is disabled in open-source // builds because it can lead to performance regressions or crashes due to // interactions with external LAPACK/BLAS implementations (e.g., // OpenBLAS/MKL). run_chunk(0, batch_count); return ffi::Error::Success(); #endif // PLATFORM_GOOGLE } //== Triangular System Solver ==// template ffi::Error TriMatrixEquationSolver::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer y, ffi::ResultBuffer y_out, MatrixParams::Side side, MatrixParams::UpLo uplo, MatrixParams::Transpose trans_x, MatrixParams::Diag diag) { CopyIfDiffBuffer(y, y_out); FFI_ASSIGN_OR_RETURN((auto [batch_count, y_rows, y_cols]), SplitBatch2D(y.dimensions())); auto* y_out_data = y_out->typed_data(); IntType x_leading_dim_v = side == MatrixParams::Side::kLeft ? y_rows : y_cols; IntType y_leading_dim_v = y_rows; auto side_v = static_cast(side); auto uplo_v = static_cast(uplo); auto trans_x_v = static_cast(trans_x); auto diag_v = static_cast(diag); FFI_ASSIGN_OR_RETURN(auto y_rows_v, MaybeCastNoOverflow(y_rows)); FFI_ASSIGN_OR_RETURN(auto y_cols_v, MaybeCastNoOverflow(y_cols)); auto* x_data = x.typed_data(); const int64_t y_out_step{y_rows * y_cols}; const int64_t x_step{x_leading_dim_v * x_leading_dim_v}; ffi::NativeType alpha = static_cast>(1); int64_t cost_per_matrix = y_rows * y_cols * (side == MatrixParams::Side::kLeft ? y_rows : y_cols); if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_y_out_data = y_out_data + i * y_out_step; const auto* local_x_data = x_data + i * x_step; for (int64_t j = 0; j < count; ++j) { TriMatrixEquationSolver::fn( &side_v, &uplo_v, &trans_x_v, &diag_v, &y_rows_v, &y_cols_v, &alpha, const_cast*>(local_x_data), &x_leading_dim_v, local_y_out_data, &y_leading_dim_v); local_y_out_data += y_out_step; local_x_data += x_step; } }); } template ffi::Error TriMatrixEquationSolverKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer y, ffi::ResultBuffer y_out, MatrixParams::Side side, MatrixParams::UpLo uplo, MatrixParams::Transpose trans_x, MatrixParams::Diag diag) { if (TriMatrixEquationSolver::fn != nullptr) { return TriMatrixEquationSolver::Kernel( thread_pool, x, y, y_out, side, uplo, trans_x, diag); } return TriMatrixEquationSolver::Kernel( thread_pool, x, y, y_out, side, uplo, trans_x, diag); } template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; template struct TriMatrixEquationSolver; //== LU Decomposition ==// template ffi::Error LuDecomposition::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer ipiv, ffi::ResultBuffer info) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* ipiv_data = ipiv->typed_data(); auto* info_data = info->typed_data(); CopyIfDiffBuffer(x, x_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); auto x_leading_dim_v = x_rows_v; const int64_t x_out_step{x_rows * x_cols}; const int64_t ipiv_step{std::min(x_rows, x_cols)}; int64_t cost_per_matrix = x_rows * x_cols * std::min(x_rows, x_cols); if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_ipiv_data = ipiv_data + i * ipiv_step; auto* local_info_data = info_data + i; std::unique_ptr ipiv_tmp; if constexpr (!std::is_same_v) { ipiv_tmp = std::make_unique(ipiv_step); } for (int64_t j = 0; j < count; ++j) { IntType info_v; if constexpr (std::is_same_v) { LuDecomposition::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_ipiv_data, &info_v); } else { LuDecomposition::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, ipiv_tmp.get(), &info_v); for (int64_t k = 0; k < ipiv_step; ++k) { local_ipiv_data[k] = static_cast(ipiv_tmp[k]); } } ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(&info_v, sizeof(info_v)); *local_info_data = static_cast(info_v); using T = ffi::NativeType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_x_out_data, x_rows * x_cols * sizeof(T)); if constexpr (std::is_same_v) { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_ipiv_data, ipiv_step * sizeof(int32_t)); } else { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(ipiv_tmp.get(), ipiv_step * sizeof(IntType)); } local_x_out_data += x_out_step; local_ipiv_data += ipiv_step; ++local_info_data; } }); } template ffi::Error LuDecompositionKernel(ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer ipiv, ffi::ResultBuffer info) { if (LuDecomposition::fn != nullptr) { return LuDecomposition::Kernel(thread_pool, x, x_out, ipiv, info); } return LuDecomposition::Kernel(thread_pool, x, x_out, ipiv, info); } template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; template struct LuDecomposition; //== QR Factorization ==// template ffi::Error QrFactorization::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer tau) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* tau_data = tau->typed_data(); CopyIfDiffBuffer(x, x_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); auto x_leading_dim_v = x_rows_v; const int64_t x_out_step{x_rows * x_cols}; const int64_t tau_step{std::min(x_rows, x_cols)}; const int64_t work_size = QrFactorization::GetWorkspaceSize(x_rows, x_cols); FFI_ASSIGN_OR_RETURN(auto workspace_dim_v, MaybeCastNoOverflow(work_size)); int64_t min_dim = std::min(x_rows, x_cols); int64_t cost_per_matrix = 2 * x_rows * x_cols * min_dim - (2 * min_dim * min_dim * min_dim) / 3; if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_tau_data = tau_data + i * tau_step; auto work_data = AllocateScratchMemory(work_size); for (int64_t j = 0; j < count; ++j) { IntType info; QrFactorization::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_tau_data, work_data.get(), &workspace_dim_v, &info); using T = ffi::NativeType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_x_out_data, x_rows * x_cols * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_tau_data, std::min(x_rows, x_cols) * sizeof(T)); local_x_out_data += x_out_step; local_tau_data += tau_step; } }); } template ffi::Error QrFactorizationKernel(ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer tau) { if (QrFactorization::fn != nullptr) { return QrFactorization::Kernel(thread_pool, x, x_out, tau); } return QrFactorization::Kernel(thread_pool, x, x_out, tau); } template int64_t QrFactorization::GetWorkspaceSize(IntType x_rows, IntType x_cols) { ValueType optimal_size{}; IntType x_leading_dim_v = x_rows; IntType info = 0; IntType workspace_query = -1; fn(&x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, &optimal_size, &workspace_query, &info); return info == 0 ? static_cast(std::real(optimal_size)) : -1; } template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; template struct QrFactorization; //== Column Pivoting QR Factorization ==// // lapack geqp3 template ffi::Error PivotingQrFactorization::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer jpvt, ffi::ResultBuffer x_out, ffi::ResultBuffer jpvt_out, ffi::ResultBuffer tau) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* jpvt_out_data = jpvt_out->typed_data(); auto* tau_data = tau->typed_data(); CopyIfDiffBuffer(x, x_out); CopyIfDiffBuffer(jpvt, jpvt_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); auto x_leading_dim_v = x_rows_v; const int64_t x_out_step{x_rows * x_cols}; const int64_t jpvt_step{x_cols}; const int64_t tau_step{std::min(x_rows, x_cols)}; const int64_t work_size = PivotingQrFactorization::GetWorkspaceSize(x_rows, x_cols); FFI_ASSIGN_OR_RETURN(auto workspace_dim_v, MaybeCastNoOverflow(work_size)); constexpr bool is_complex_dtype = ffi::IsComplexType(); int64_t min_dim = std::min(x_rows, x_cols); int64_t cost_per_matrix = 2 * x_rows * x_cols * min_dim - (2 * min_dim * min_dim * min_dim) / 3; if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_jpvt_out_data = jpvt_out_data + i * jpvt_step; auto* local_tau_data = tau_data + i * tau_step; auto work_data = AllocateScratchMemory(work_size); std::unique_ptr[]> rwork_data; if constexpr (is_complex_dtype) { rwork_data = AllocateScratchMemory(2 * x_cols); } std::unique_ptr jpvt_tmp; if constexpr (!std::is_same_v) { jpvt_tmp = std::make_unique(jpvt_step); } for (int64_t j = 0; j < count; ++j) { IntType info_v; if constexpr (std::is_same_v) { if constexpr (is_complex_dtype) { PivotingQrFactorization::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_jpvt_out_data, local_tau_data, work_data.get(), &workspace_dim_v, rwork_data.get(), &info_v); } else { PivotingQrFactorization::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_jpvt_out_data, local_tau_data, work_data.get(), &workspace_dim_v, &info_v); } ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_jpvt_out_data, jpvt_step * sizeof(int32_t)); using T = ffi::NativeType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_x_out_data, x_rows * x_cols * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_tau_data, std::min(x_rows, x_cols) * sizeof(T)); } else { for (int64_t k = 0; k < jpvt_step; ++k) { jpvt_tmp[k] = static_cast(local_jpvt_out_data[k]); } if constexpr (is_complex_dtype) { PivotingQrFactorization::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, jpvt_tmp.get(), local_tau_data, work_data.get(), &workspace_dim_v, rwork_data.get(), &info_v); } else { PivotingQrFactorization::fn( &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, jpvt_tmp.get(), local_tau_data, work_data.get(), &workspace_dim_v, &info_v); } ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(jpvt_tmp.get(), jpvt_step * sizeof(IntType)); using T = ffi::NativeType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_x_out_data, x_rows * x_cols * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_tau_data, std::min(x_rows, x_cols) * sizeof(T)); for (int64_t k = 0; k < jpvt_step; ++k) { local_jpvt_out_data[k] = static_cast(jpvt_tmp[k]); } } local_x_out_data += x_out_step; local_jpvt_out_data += jpvt_step; local_tau_data += tau_step; } }); } template ffi::Error PivotingQrFactorizationKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer jpvt, ffi::ResultBuffer x_out, ffi::ResultBuffer jpvt_out, ffi::ResultBuffer tau) { if (PivotingQrFactorization::fn != nullptr) { return PivotingQrFactorization::Kernel( thread_pool, x, jpvt, x_out, jpvt_out, tau); } return PivotingQrFactorization::Kernel(thread_pool, x, jpvt, x_out, jpvt_out, tau); } template int64_t PivotingQrFactorization::GetWorkspaceSize( IntType x_rows, IntType x_cols) { ValueType optimal_size{}; IntType x_leading_dim_v = x_rows; IntType info = 0; IntType workspace_query = -1; if constexpr (ffi::IsComplexType()) { fn(&x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &optimal_size, &workspace_query, nullptr, &info); } else { fn(&x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &optimal_size, &workspace_query, &info); } return info == 0 ? static_cast(std::real(optimal_size)) : -1; } template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; template struct PivotingQrFactorization; //== Orthogonal QR ==// //== Computes orthogonal matrix Q from QR Decomposition ==// template ffi::Error OrthogonalQr::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer tau, ffi::ResultBuffer x_out) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* tau_data = tau.typed_data(); auto* x_out_data = x_out->typed_data(); CopyIfDiffBuffer(x, x_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); FFI_ASSIGN_OR_RETURN(auto tau_size_v, MaybeCastNoOverflow(tau.dimensions().back())); auto x_leading_dim_v = x_rows_v; const int64_t x_out_step{x_rows * x_cols}; const int64_t tau_step{tau_size_v}; int64_t work_size = OrthogonalQr::GetWorkspaceSize( x_rows, x_cols, tau.dimensions().back()); FFI_ASSIGN_OR_RETURN(auto work_size_v, MaybeCastNoOverflow(work_size)); int64_t k_dim = tau_size_v; int64_t cost_per_matrix = 4 * x_rows * x_cols * k_dim - 2 * k_dim * k_dim * (x_rows + x_cols) + (4 * k_dim * k_dim * k_dim) / 3; if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_tau_data = tau_data + i * tau_step; auto work_data = AllocateScratchMemory(work_size); for (int64_t j = 0; j < count; ++j) { IntType info; OrthogonalQr::fn(&x_rows_v, &x_cols_v, &tau_size_v, local_x_out_data, &x_leading_dim_v, local_tau_data, work_data.get(), &work_size_v, &info); local_x_out_data += x_out_step; local_tau_data += tau_step; } }); } template ffi::Error OrthogonalQrKernel(ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::Buffer tau, ffi::ResultBuffer x_out) { if (OrthogonalQr::fn != nullptr) { return OrthogonalQr::Kernel(thread_pool, x, tau, x_out); } return OrthogonalQr::Kernel(thread_pool, x, tau, x_out); } template int64_t OrthogonalQr::GetWorkspaceSize(IntType x_rows, IntType x_cols, IntType tau_size) { ValueType optimal_size = {}; IntType x_leading_dim_v = x_rows; IntType info = 0; IntType workspace_query = -1; fn(&x_rows, &x_cols, &tau_size, nullptr, &x_leading_dim_v, nullptr, &optimal_size, &workspace_query, &info); return info == 0 ? static_cast(std::real(optimal_size)) : -1; } template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; template struct OrthogonalQr; //== Orthogonal QR Multiply ==// template ffi::Error OrthogonalQrMultiply::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer a, ffi::Buffer tau, ffi::Buffer c, bool left, bool transpose, ffi::ResultBuffer c_out) { FFI_ASSIGN_OR_RETURN((auto [batch_count, c_rows, c_cols]), SplitBatch2D(c.dimensions())); FFI_ASSIGN_OR_RETURN((auto [a_batch, a_rows, a_cols]), SplitBatch2D(a.dimensions())); auto* tau_data = tau.typed_data(); auto* a_data = a.typed_data(); auto* c_out_data = c_out->typed_data(); CopyIfDiffBuffer(c, c_out); char side_v = left ? 'L' : 'R'; char trans_v; if constexpr (ffi::IsComplexType()) { trans_v = transpose ? 'C' : 'N'; } else { trans_v = transpose ? 'T' : 'N'; } FFI_ASSIGN_OR_RETURN(auto c_rows_v, MaybeCastNoOverflow(c_rows)); FFI_ASSIGN_OR_RETURN(auto c_cols_v, MaybeCastNoOverflow(c_cols)); FFI_ASSIGN_OR_RETURN(auto k_v, MaybeCastNoOverflow(tau.dimensions().back())); // LDA is the leading dimension of A: m (= c_rows) when left, n (= c_cols) // when right. The shape rule guarantees a_rows equals the correct value in // both cases, but we compute it explicitly to match the trsm convention. FFI_ASSIGN_OR_RETURN(auto lda_v, MaybeCastNoOverflow(left ? a_rows : c_cols)); int64_t work_size = GetWorkspaceSize(side_v, trans_v, c_rows_v, c_cols_v, k_v, lda_v); FFI_ASSIGN_OR_RETURN(auto work_size_v, MaybeCastNoOverflow(work_size)); auto c_leading_dim_v = c_rows_v; const int64_t c_out_step{c_rows * c_cols}; const int64_t a_step{a_rows * a_cols}; const int64_t tau_step{tau.dimensions().back()}; int64_t k_dim = k_v; int64_t cost_per_matrix = 4 * k_dim * c_rows * c_cols - 2 * k_dim * k_dim * (left ? c_cols : c_rows); if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_c_out_data = c_out_data + i * c_out_step; const auto* local_a_data = a_data + i * a_step; const auto* local_tau_data = tau_data + i * tau_step; auto work_data = AllocateScratchMemory(work_size); for (int64_t j = 0; j < count; ++j) { IntType info; fn(&side_v, &trans_v, &c_rows_v, &c_cols_v, &k_v, const_cast(local_a_data), &lda_v, const_cast(local_tau_data), local_c_out_data, &c_leading_dim_v, work_data.get(), &work_size_v, &info); local_c_out_data += c_out_step; local_a_data += a_step; local_tau_data += tau_step; } }); } template int64_t OrthogonalQrMultiply::GetWorkspaceSize( char side, char trans, IntType m, IntType n, IntType k, IntType lda) { ValueType optimal_size = {}; IntType c_leading_dim = m; IntType info = 0; IntType workspace_query = -1; fn(&side, &trans, &m, &n, &k, nullptr, &lda, nullptr, nullptr, &c_leading_dim, &optimal_size, &workspace_query, &info); return info == 0 ? static_cast(std::real(optimal_size)) : -1; } template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template struct OrthogonalQrMultiply; template ffi::Error OrthogonalQrMultiplyKernel(ffi::ThreadPool thread_pool, ffi::Buffer a, ffi::Buffer tau, ffi::Buffer c, bool left, bool transpose, ffi::ResultBuffer c_out) { if (OrthogonalQrMultiply::fn != nullptr) { return OrthogonalQrMultiply::Kernel(thread_pool, a, tau, c, left, transpose, c_out); } return OrthogonalQrMultiply::Kernel(thread_pool, a, tau, c, left, transpose, c_out); } //== Cholesky Factorization ==// template ffi::Error CholeskyFactorization::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, MatrixParams::UpLo uplo, ffi::ResultBuffer x_out, ffi::ResultBuffer info) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* info_data = info->typed_data(); CopyIfDiffBuffer(x, x_out); auto uplo_v = static_cast(uplo); FFI_ASSIGN_OR_RETURN(auto n_v, MaybeCastNoOverflow(x_rows)); auto lda_v = n_v; const int64_t x_out_step{x_rows * x_cols}; int64_t cost_per_matrix = x_rows * x_rows * x_rows / 3; if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_info_data = info_data + i; for (int64_t j = 0; j < count; ++j) { IntType info_v; CholeskyFactorization::fn( &uplo_v, &n_v, local_x_out_data, &lda_v, &info_v); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(&info_v, sizeof(info_v)); *local_info_data = static_cast(info_v); using T = ffi::NativeType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(local_x_out_data, x_rows * x_cols * sizeof(T)); local_x_out_data += x_out_step; ++local_info_data; } }); } template ffi::Error CholeskyFactorizationKernel(ffi::ThreadPool thread_pool, ffi::Buffer x, MatrixParams::UpLo uplo, ffi::ResultBuffer x_out, ffi::ResultBuffer info) { if (CholeskyFactorization::fn != nullptr) { return CholeskyFactorization::Kernel(thread_pool, x, uplo, x_out, info); } return CholeskyFactorization::Kernel(thread_pool, x, uplo, x_out, info); } template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; template struct CholeskyFactorization; //== Singular Value Decomposition (SVD) ==// //== using a divide and conquer method ==// namespace internal { // TODO(basioli): It might be worth consider parallelizing on even smaller // matrices if batch size is large enough. // In this case we wouldn't queue a call per matrix, but a call where we'd // handle a chunk of the batch (e.x. (10000, 2, 2) could be chunked into e.x. // (250, 2, 2)). bool ShouldParallelizeSVD(int64_t batch_size, int64_t rows, int64_t cols, int64_t num_threads) { int64_t matrix_size = rows * cols; const int64_t kMinMatrixSizeForParallelization = 8 * 8; return matrix_size >= kMinMatrixSizeForParallelization && batch_size > 1; } template static ffi::Error SvdKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (mode == svd::ComputationMode::kComputeVtOverwriteXPartialU) [[unlikely]] { return ffi::Error( XLA_FFI_Error_Code_UNIMPLEMENTED, "Current implementation does not support this computation mode"); } FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* singular_values_data = singular_values->typed_data(); auto* u_data = u->typed_data(); auto* vt_data = vt->typed_data(); auto* info_data = info->typed_data(); CopyIfDiffBuffer(x, x_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); auto mode_v = static_cast(mode); auto x_leading_dim_v = x_rows_v; auto u_leading_dim_v = x_rows_v; auto u_dims = u->dimensions().last(2); auto vt_dims = vt->dimensions().last(2); FFI_ASSIGN_OR_RETURN(auto vt_leading_dim_v, MaybeCastNoOverflow(vt_dims.front())); // Prepare LAPACK workspaces. auto work_size_or = svd::SVDType::GetWorkspaceSize(x_rows_v, x_cols_v, mode); FFI_ASSIGN_OR_RETURN(const auto work_size, work_size_or); const auto iwork_size = svd::GetIntWorkspaceSize(x_rows, x_cols); FFI_ASSIGN_OR_RETURN(auto workspace_dim_v, MaybeCastNoOverflow(work_size)); int64_t rwork_size = 0; if constexpr (ffi::IsComplexType()) { rwork_size = svd::GetRealWorkspaceSize(x_rows, x_cols, mode); } const int64_t x_out_step{x_rows * x_cols}; const int64_t singular_values_step{singular_values->dimensions().back()}; const int64_t u_step{u_dims.front() * u_dims.back()}; const int64_t vt_step{vt_leading_dim_v * vt_dims.back()}; int64_t min_dim = std::min(x_rows, x_cols); int64_t cost_per_matrix = 10 * x_rows * x_cols * min_dim; int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_singular_values_data = singular_values_data + i * singular_values_step; auto* local_u_data = u_data + i * u_step; auto* local_vt_data = vt_data + i * vt_step; auto* local_info_data = info_data + i; auto work_data = AllocateScratchMemory(work_size); auto iwork_data = AllocateScratchMemory>(iwork_size); std::unique_ptr[]> rwork; if constexpr (ffi::IsComplexType()) { rwork = AllocateScratchMemory(rwork_size); } for (int64_t j = 0; j < count; ++j) { IntType info_v; if constexpr (ffi::IsComplexType()) { svd::SVDType::fn( &mode_v, &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_singular_values_data, local_u_data, &u_leading_dim_v, local_vt_data, &vt_leading_dim_v, work_data.get(), &workspace_dim_v, rwork.get(), iwork_data.get(), &info_v); } else { svd::SVDType::fn( &mode_v, &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_singular_values_data, local_u_data, &u_leading_dim_v, local_vt_data, &vt_leading_dim_v, work_data.get(), &workspace_dim_v, iwork_data.get(), &info_v); } // Suppress MSAN warnings when using a copy of LAPACK uninstrumented // by MSAN. using T [[maybe_unused]] = typename svd::SVDType::ValueType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(&info_v, sizeof(info_v)); *local_info_data = static_cast(info_v); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_x_out_data, x_cols_v * x_leading_dim_v * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_singular_values_data, std::min(x_rows_v, x_cols_v) * sizeof(ffi::NativeType)); if (mode_v == 'A') { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * x_rows_v * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } else if (mode_v == 'O') { if (x_rows_v < x_cols_v) { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * x_rows_v * sizeof(T)); } else { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } } else if (mode_v == 'S') { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * std::min(x_rows_v, x_cols_v) * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } local_x_out_data += x_out_step; local_singular_values_data += singular_values_step; local_u_data += u_step; local_vt_data += vt_step; ++local_info_data; } }); } template static int64_t SvdGetWorkspaceSize(IntType x_rows, IntType x_cols, svd::ComputationMode mode) { ffi::NativeType optimal_size = {}; IntType info = 0; IntType workspace_query = -1; auto mode_v = static_cast(mode); auto x_leading_dim_v = x_rows; auto u_leading_dim_v = x_rows; auto vt_leading_dim_v = mode == svd::ComputationMode::kComputeFullUVt ? x_cols : std::min(x_rows, x_cols); if constexpr (ffi::IsComplexType()) { svd::SVDType::fn( &mode_v, &x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &u_leading_dim_v, nullptr, &vt_leading_dim_v, &optimal_size, &workspace_query, nullptr, nullptr, &info); } else { svd::SVDType::fn( &mode_v, &x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &u_leading_dim_v, nullptr, &vt_leading_dim_v, &optimal_size, &workspace_query, nullptr, &info); } return info == 0 ? static_cast(std::real(optimal_size)) : -1; } template static ffi::Error SvdQRKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (mode == svd::ComputationMode::kComputeVtOverwriteXPartialU) [[unlikely]] { return ffi::Error( XLA_FFI_Error_Code_UNIMPLEMENTED, "SVD: Current implementation does not support this computation mode"); } FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* singular_values_data = singular_values->typed_data(); auto* u_data = u->typed_data(); auto* vt_data = vt->typed_data(); auto* info_data = info->typed_data(); CopyIfDiffBuffer(x, x_out); FFI_ASSIGN_OR_RETURN(auto x_rows_v, MaybeCastNoOverflow(x_rows)); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); auto mode_v = static_cast(mode); auto x_leading_dim_v = x_rows_v; auto u_leading_dim_v = x_rows_v; // Prepare LAPACK workspaces. auto work_size_or = svd::SVDQRType::GetWorkspaceSize( x_rows_v, x_cols_v, mode); FFI_ASSIGN_OR_RETURN(auto work_size, work_size_or); auto workspace_dim_v = work_size; auto u_dims = u->dimensions().last(2); auto vt_dims = vt->dimensions().last(2); FFI_ASSIGN_OR_RETURN(auto vt_leading_dim_v, MaybeCastNoOverflow(vt_dims.front())); const int64_t x_out_step{x_rows * x_cols}; const int64_t singular_values_step{singular_values->dimensions().back()}; const int64_t u_step{u_dims.front() * u_dims.back()}; const int64_t vt_step{vt_leading_dim_v * vt_dims.back()}; int64_t cost_per_matrix = 10 * x_rows * x_cols * std::min(x_rows, x_cols); if constexpr (ffi::IsComplexType()) { cost_per_matrix *= 3; } int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_singular_values_data = singular_values_data + i * singular_values_step; auto* local_u_data = u_data + i * u_step; auto* local_vt_data = vt_data + i * vt_step; auto* local_info_data = info_data + i; auto work_data = AllocateScratchMemory(work_size); std::unique_ptr[]> rwork; if constexpr (ffi::IsComplexType()) { const auto rwork_size = svd::GetRealWorkspaceSizeQR(x_rows, x_cols); rwork = AllocateScratchMemory(rwork_size); } for (int64_t j = 0; j < count; ++j) { IntType info_v; if constexpr (ffi::IsComplexType()) { svd::SVDQRType::fn( &mode_v, &mode_v, &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_singular_values_data, local_u_data, &u_leading_dim_v, local_vt_data, &vt_leading_dim_v, work_data.get(), &workspace_dim_v, rwork.get(), &info_v); } else { svd::SVDQRType::fn( &mode_v, &mode_v, &x_rows_v, &x_cols_v, local_x_out_data, &x_leading_dim_v, local_singular_values_data, local_u_data, &u_leading_dim_v, local_vt_data, &vt_leading_dim_v, work_data.get(), &workspace_dim_v, &info_v); } using T [[maybe_unused]] = typename svd::SVDQRType::ValueType; ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(&info_v, sizeof(info_v)); *local_info_data = static_cast(info_v); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_x_out_data, x_cols_v * x_leading_dim_v * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_singular_values_data, std::min(x_rows_v, x_cols_v) * sizeof(ffi::NativeType)); if (mode_v == 'A') { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * x_rows_v * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } else if (mode_v == 'O') { if (x_rows_v < x_cols_v) { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * x_rows_v * sizeof(T)); } else { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } } else if (mode_v == 'S') { ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_u_data, u_leading_dim_v * std::min(x_rows_v, x_cols_v) * sizeof(T)); ABSL_ANNOTATE_MEMORY_IS_INITIALIZED( local_vt_data, vt_leading_dim_v * x_cols_v * sizeof(T)); } local_x_out_data += x_out_step; local_singular_values_data += singular_values_step; local_u_data += u_step; local_vt_data += vt_step; ++local_info_data; } }); } template static absl::StatusOr SvdQRGetWorkspaceSize( IntType x_rows, IntType x_cols, svd::ComputationMode mode) { ffi::NativeType optimal_size = {}; IntType info = 0; IntType workspace_query = -1; auto mode_v = static_cast(mode); auto x_leading_dim_v = x_rows; auto u_leading_dim_v = x_rows; auto vt_leading_dim_v = mode == svd::ComputationMode::kComputeFullUVt ? x_cols : std::min(x_rows, x_cols); if constexpr (ffi::IsComplexType()) { svd::SVDQRType::fn( &mode_v, &mode_v, &x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &u_leading_dim_v, nullptr, &vt_leading_dim_v, &optimal_size, &workspace_query, nullptr, &info); } else { svd::SVDQRType::fn( &mode_v, &mode_v, &x_rows, &x_cols, nullptr, &x_leading_dim_v, nullptr, nullptr, &u_leading_dim_v, nullptr, &vt_leading_dim_v, &optimal_size, &workspace_query, &info); } return info == 0 ? MaybeCastNoOverflow(std::real(optimal_size)) : -1; } } // namespace internal template ffi::Error SingularValueDecompositionKernel( ::xla::ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (svd::SVDType::fn != nullptr) { return internal::SvdKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } return internal::SvdKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } template ffi::Error SingularValueDecompositionComplexKernel( ::xla::ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (svd::SVDType::fn != nullptr) { return internal::SvdKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } return internal::SvdKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } template absl::StatusOr SingularValueDecomposition::GetWorkspaceSize( IntType x_rows, IntType x_cols, svd::ComputationMode mode) { return internal::SvdGetWorkspaceSize(x_rows, x_cols, mode); } template absl::StatusOr SingularValueDecompositionComplex::GetWorkspaceSize( IntType x_rows, IntType x_cols, svd::ComputationMode mode) { return internal::SvdGetWorkspaceSize(x_rows, x_cols, mode); } template ffi::Error SingularValueDecompositionQRKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (svd::SVDQRType::fn != nullptr) { return internal::SvdQRKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } return internal::SvdQRKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } template ffi::Error SingularValueDecompositionQRComplexKernel( ffi::ThreadPool thread_pool, ffi::Buffer x, ffi::ResultBuffer x_out, ffi::ResultBuffer singular_values, ffi::ResultBuffer u, ffi::ResultBuffer vt, ffi::ResultBuffer info, svd::ComputationMode mode) { if (svd::SVDQRType::fn != nullptr) { return internal::SvdQRKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } return internal::SvdQRKernel( thread_pool, x, x_out, singular_values, u, vt, info, mode); } template absl::StatusOr SingularValueDecompositionQR::GetWorkspaceSize( IntType x_rows, IntType x_cols, svd::ComputationMode mode) { return internal::SvdQRGetWorkspaceSize(x_rows, x_cols, mode); } template absl::StatusOr SingularValueDecompositionQRComplex::GetWorkspaceSize( IntType x_rows, IntType x_cols, svd::ComputationMode mode) { return internal::SvdQRGetWorkspaceSize(x_rows, x_cols, mode); } int64_t svd::GetRealWorkspaceSize(int64_t x_rows, int64_t x_cols, svd::ComputationMode mode) { const auto min_dim = std::min(x_rows, x_cols); if (!ComputesUV(mode)) { return 7 * min_dim; } const auto max_dim = std::max(x_rows, x_cols); return std::max(5 * min_dim * min_dim + 5 * min_dim, 2 * max_dim * min_dim + 2 * min_dim * min_dim + min_dim); } int64_t svd::GetRealWorkspaceSizeQR(int64_t x_rows, int64_t x_cols) { return 5 * std::min(x_rows, x_cols); } int64_t svd::GetIntWorkspaceSize(int64_t x_rows, int64_t x_cols) { return 8 * std::min(x_rows, x_cols); } template struct SingularValueDecomposition; template struct SingularValueDecomposition; template struct SingularValueDecomposition; template struct SingularValueDecomposition; template struct SingularValueDecompositionComplex; template struct SingularValueDecompositionComplex; template struct SingularValueDecompositionComplex; template struct SingularValueDecompositionComplex; template struct SingularValueDecompositionQR; template struct SingularValueDecompositionQR; template struct SingularValueDecompositionQR; template struct SingularValueDecompositionQR; template struct SingularValueDecompositionQRComplex; template struct SingularValueDecompositionQRComplex; template struct SingularValueDecompositionQRComplex; template struct SingularValueDecompositionQRComplex; //== Eigenvalues and eigenvectors ==// int64_t eig::GetWorkspaceSize(int64_t x_cols, ComputationMode mode) { switch (mode) { case ComputationMode::kNoEigenvectors: return 2 * x_cols + 1; case ComputationMode::kComputeEigenvectors: return 1 + 6 * x_cols + 2 * x_cols * x_cols; } } int64_t eig::GetIntWorkspaceSize(int64_t x_cols, ComputationMode mode) { switch (mode) { case ComputationMode::kNoEigenvectors: return 1; case ComputationMode::kComputeEigenvectors: return 3 + 5 * x_cols; } } template ffi::Error EigenvalueDecompositionSymmetric::Kernel( ffi::ThreadPool thread_pool, ffi::Buffer x, MatrixParams::UpLo uplo, ffi::ResultBuffer x_out, ffi::ResultBuffer eigenvalues, ffi::ResultBuffer info, eig::ComputationMode mode) { FFI_ASSIGN_OR_RETURN((auto [batch_count, x_rows, x_cols]), SplitBatch2D(x.dimensions())); auto* x_out_data = x_out->typed_data(); auto* eigenvalues_data = eigenvalues->typed_data(); auto* info_data = info->typed_data(); CopyIfDiffBuffer(x, x_out); auto mode_v = static_cast(mode); auto uplo_v = static_cast(uplo); FFI_ASSIGN_OR_RETURN(auto x_cols_v, MaybeCastNoOverflow(x_cols)); FFI_ASSIGN_OR_RETURN(auto x_leading_dim_v, MaybeCastNoOverflow(x_cols)); const int64_t x_out_step{x_cols * x_cols}; const int64_t eigenvalues_step{x_cols}; const int64_t work_size = eig::GetWorkspaceSize(x_cols, mode); FFI_ASSIGN_OR_RETURN(auto work_size_v, MaybeCastNoOverflow(work_size)); const int64_t iwork_size = eig::GetIntWorkspaceSize(x_cols, mode); FFI_ASSIGN_OR_RETURN(auto iwork_size_v, MaybeCastNoOverflow(iwork_size)); int64_t cost_per_matrix = 5 * x_cols * x_cols * x_cols; int64_t chunk_size = GetLapackBatchChunkSize(batch_count, cost_per_matrix, thread_pool.num_threads()); return ParallelBatchMap( thread_pool, batch_count, chunk_size, [&](int64_t i, int64_t count) { auto* local_x_out_data = x_out_data + i * x_out_step; auto* local_eigenvalues_data = eigenvalues_data + i * eigenvalues_step; auto* local_info_data = info_data + i; auto work_data = AllocateScratchMemory