Skip to content

[WA] Fix the lack of exceptions vector in getrf_batch #1916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[WA] Fix the lack of exceptions vector in getrf_batch #1916

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1a9243a
Fix the lack of exceptions vector in getrf_batch
CuiYifeng Jul 14, 2025
29605c0
Fix typo and warning
CuiYifeng Aug 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions src/ATen/native/xpu/mkl/BatchLinearAlgebra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <ATen/ops/_linalg_check_errors_native.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/from_blob.h>
#include <ATen/ops/zeros_like.h>

#include <comm/SYCLContext.h>
#include <comm/TensorInfo.h>
Expand Down Expand Up @@ -40,6 +41,16 @@ static oneapi::mkl::transpose to_blas_(TransposeType trans) {
void error_handle(int32_t* infos, const oneapi::mkl::lapack::batch_error& be) {
auto errs = be.exceptions();
auto ids = be.ids();

if (!errs.size()) {
TORCH_WARN("Caught lapack exception:\nWhat: ", be.what(), "\nInfo: ", be.info());
for (auto& i : ids) {
TORCH_WARN("Error in matrix #", i);
infos[i] = 1;
}
return;
}

for (auto& i : ids) {
try {
std::rethrow_exception(errs[i]);
Expand Down Expand Up @@ -529,8 +540,8 @@ void lu_factor_mkl(
"linalg.lu_factor: LU without pivoting is not implemented on the XPU");

// handle the info
info.zero_();
int32_t* infos_data = info.data_ptr<int32_t>();
Tensor info_ = at::zeros_like(info, Device(at::kCPU));
int32_t* infos_data = info_.data_ptr<int32_t>();
Copy link
Contributor

@jianyizh jianyizh Aug 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you provide more explanation here? why you create a zero tensor on cpu and pass it to apply_lu_xpu_？ and then you copy info_ to xpu? mkl will get seg fault when info is on xpu?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On stock Pytorch side, original info Tensor is on device instead of host.


// oneMKL requires Long for pivots but PyTorch provides Int
Tensor pivots_ = at::empty(pivots.sizes(), pivots.options().dtype(kLong));
Expand All @@ -539,7 +550,8 @@ void lu_factor_mkl(
apply_lu_xpu_<scalar_t>(LU, pivots_, infos_data);
});

// Copy to original pivots tensor
// Copy to original info and pivots tensor
info.copy_(info_);
pivots.copy_(pivots_);
}

Expand Down