apply_lu_factor_batched_cublas Class — pytorch Architecture

Source Code

aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLibBlas.cpp lines 107–122

template <typename scalar_t>
static void apply_lu_factor_batched_cublas(const Tensor& A, const Tensor& pivots, const Tensor& infos, bool get_pivots) {
  // This function just works with square matrices
  TORCH_INTERNAL_ASSERT(A.size(-2) == A.size(-1));

  auto batch_size = cuda_int_cast(batchCount(A), "batch_size");;
  auto n = cuda_int_cast(A.size(-2), "n");
  auto lda = cuda_int_cast(std::max<int>(1, n), "lda");

  auto pivots_data = get_pivots ? pivots.data_ptr<int>() : nullptr;
  auto infos_data = infos.data_ptr<int>();
  Tensor a_ptr_array = get_device_pointers<scalar_t>(A);
  auto a_ptr_array_data = reinterpret_cast<scalar_t**>(a_ptr_array.data_ptr());

  at::cuda::blas::getrfBatched(n, a_ptr_array_data, lda, pivots_data, infos_data, batch_size);
}