From 97490b547d8ec5c5de943041ac6d697fe3cb1b18 Mon Sep 17 00:00:00 2001
From: Kerelos Gayed <kerelos.gayed@shopify.com>
Date: Fri, 25 Sep 2026 19:44:05 -0400
Subject: [PATCH 1/2] Fix CUDA Cholesky workspace handling

---
 xpu/linear_algebra.hpp | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/xpu/linear_algebra.hpp b/xpu/linear_algebra.hpp
index 87426db..09e6c01 100644
--- a/xpu/linear_algebra.hpp
+++ b/xpu/linear_algebra.hpp
@@ -191,20 +191,19 @@ inline auto cusolver_potrf(
   std::size_t order,
   std::size_t stride,
   T* matrix,
-  T* workspace,
-  std::size_t workspace_size,
+  xpu::buffer<T>& workspace,
   int* info
 ) -> void {
   const auto vendor_order{xpu::detail::checked_cast<int>(order)};
   const auto vendor_stride{xpu::detail::checked_cast<int>(stride)};
-  const auto vendor_workspace_size{xpu::detail::checked_cast<int>(workspace_size)};
+  const auto vendor_workspace_size{xpu::detail::checked_cast<int>(workspace.count())};
 
   if constexpr (std::same_as<T, float>) {
     xpu::cu_check(cusolverDnSpotrf(
       handle, CUBLAS_FILL_MODE_UPPER,
       vendor_order,
       matrix, vendor_stride,
-      workspace, vendor_workspace_size,
+      workspace.data(), vendor_workspace_size,
       info
     ));
   } else {
@@ -212,7 +211,7 @@ inline auto cusolver_potrf(
       handle, CUBLAS_FILL_MODE_UPPER,
       vendor_order,
       matrix, vendor_stride,
-      workspace, vendor_workspace_size,
+      workspace.data(), vendor_workspace_size,
       info
     ));
   }
@@ -555,7 +554,7 @@ class cholesky_factorization {
 #if defined(XPU_CUDA)
     detail::cusolver_potrf(
       handle_, order_, stride_,
-      matrix, workspace_.data(), workspace_.size(), info_.data()
+      matrix, workspace_, info_.data()
     );
 
     auto info{0};

From 759669c4ee9ef5325f2f8399929f26fe7fde529e Mon Sep 17 00:00:00 2001
From: Kerelos Gayed <kerelos.gayed@shopify.com>
Date: Fri, 25 Sep 2026 20:23:30 -0400
Subject: [PATCH 2/2] Keep Cholesky workspace fix at the call site

---
 xpu/linear_algebra.hpp | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/xpu/linear_algebra.hpp b/xpu/linear_algebra.hpp
index 09e6c01..273113f 100644
--- a/xpu/linear_algebra.hpp
+++ b/xpu/linear_algebra.hpp
@@ -191,19 +191,20 @@ inline auto cusolver_potrf(
   std::size_t order,
   std::size_t stride,
   T* matrix,
-  xpu::buffer<T>& workspace,
+  T* workspace,
+  std::size_t workspace_size,
   int* info
 ) -> void {
   const auto vendor_order{xpu::detail::checked_cast<int>(order)};
   const auto vendor_stride{xpu::detail::checked_cast<int>(stride)};
-  const auto vendor_workspace_size{xpu::detail::checked_cast<int>(workspace.count())};
+  const auto vendor_workspace_size{xpu::detail::checked_cast<int>(workspace_size)};
 
   if constexpr (std::same_as<T, float>) {
     xpu::cu_check(cusolverDnSpotrf(
       handle, CUBLAS_FILL_MODE_UPPER,
       vendor_order,
       matrix, vendor_stride,
-      workspace.data(), vendor_workspace_size,
+      workspace, vendor_workspace_size,
       info
     ));
   } else {
@@ -211,7 +212,7 @@ inline auto cusolver_potrf(
       handle, CUBLAS_FILL_MODE_UPPER,
       vendor_order,
       matrix, vendor_stride,
-      workspace.data(), vendor_workspace_size,
+      workspace, vendor_workspace_size,
       info
     ));
   }
@@ -554,7 +555,7 @@ class cholesky_factorization {
 #if defined(XPU_CUDA)
     detail::cusolver_potrf(
       handle_, order_, stride_,
-      matrix, workspace_, info_.data()
+      matrix, workspace_.data(), workspace_.count(), info_.data()
     );
 
     auto info{0};