Skip to content

fix(linalg): make GPU scalar in-place ops correct on CUDA (#988) #992

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

Merged
pcchen merged 2 commits into from
Jul 9, 2026
Merged

fix(linalg): make GPU scalar in-place ops correct on CUDA (#988) #992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4ac2837
fix(linalg): make GPU scalar in-place ops correct on CUDA (#988)
yingjerkao Jul 6, 2026
3fda11f
ci: retrigger checks after retarget to master
pcchen Jul 8, 2026
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
110 changes: 55 additions & 55 deletions src/Tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -547,20 +547,20 @@ namespace cytnx {
}

namespace {
// Wrap a scalar as a shape-{1} Tensor on `device` so scalar in-place
// arithmetic reuses the iAdd/iSub/iMul/iDiv kernels (which broadcast a
// length-1 RHS and mutate LHS storage in place). See #906.
// Wrap a scalar as a host-resident shape-{1} Tensor so scalar in-place arithmetic reuses the
// iAdd/iSub/iMul/iDiv kernels (which broadcast a length-1 RHS and mutate LHS storage in
// place). See #906. The wrapper deliberately stays on the CPU even when the LHS is on a GPU:
// for a length-1 RHS the GPU kernels read the scalar with a host-side dereference and pass it
// in by value, so keeping it on the host avoids a per-call H2D copy. See #988.
template <class T>
Tensor _scalar_as_rank1_tensor(const T &rc, const int device) {
Tensor _scalar_as_rank1_tensor(const T &rc) {
Tensor s({1}, Type.cy_typeid(rc), Device.cpu);
s.storage().at<T>(0) = rc;
if (device != Device.cpu) s = s.to(device);
return s;
}
Tensor _scalar_as_rank1_tensor(const Scalar &rc, const int device) {
Tensor _scalar_as_rank1_tensor(const Scalar &rc) {
Tensor s({1}, rc.dtype(), Device.cpu);
s.item() = rc; // Sproxy assignment
if (device != Device.cpu) s = s.to(device);
return s;
}
} // namespace
Expand All @@ -579,62 +579,62 @@ namespace cytnx {
}
template <>
Tensor &Tensor::operator+=<cytnx_complex128>(const cytnx_complex128 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_complex64>(const cytnx_complex64 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_double>(const cytnx_double &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_float>(const cytnx_float &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_int64>(const cytnx_int64 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_uint64>(const cytnx_uint64 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_int32>(const cytnx_int32 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_uint32>(const cytnx_uint32 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_int16>(const cytnx_int16 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_uint16>(const cytnx_uint16 &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<cytnx_bool>(const cytnx_bool &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator+=<Scalar>(const Scalar &rc) {
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iAdd(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Expand All @@ -654,62 +654,62 @@ namespace cytnx {
}
template <>
Tensor &Tensor::operator-=<cytnx_complex128>(const cytnx_complex128 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_complex64>(const cytnx_complex64 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_double>(const cytnx_double &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_float>(const cytnx_float &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_int64>(const cytnx_int64 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_uint64>(const cytnx_uint64 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_int32>(const cytnx_int32 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_uint32>(const cytnx_uint32 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_int16>(const cytnx_int16 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_uint16>(const cytnx_uint16 &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<cytnx_bool>(const cytnx_bool &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator-=<Scalar>(const Scalar &rc) {
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iSub(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Expand All @@ -729,62 +729,62 @@ namespace cytnx {
}
template <>
Tensor &Tensor::operator*=<cytnx_complex128>(const cytnx_complex128 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_complex64>(const cytnx_complex64 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_double>(const cytnx_double &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_float>(const cytnx_float &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_int64>(const cytnx_int64 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_uint64>(const cytnx_uint64 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_int32>(const cytnx_int32 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_uint32>(const cytnx_uint32 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_int16>(const cytnx_int16 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_uint16>(const cytnx_uint16 &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<cytnx_bool>(const cytnx_bool &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator*=<Scalar>(const Scalar &rc) {
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iMul(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Expand All @@ -805,62 +805,62 @@ namespace cytnx {
}
template <>
Tensor &Tensor::operator/=<cytnx_complex128>(const cytnx_complex128 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_complex64>(const cytnx_complex64 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_double>(const cytnx_double &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_float>(const cytnx_float &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_int64>(const cytnx_int64 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_uint64>(const cytnx_uint64 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_int32>(const cytnx_int32 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_uint32>(const cytnx_uint32 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_int16>(const cytnx_int16 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_uint16>(const cytnx_uint16 &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<cytnx_bool>(const cytnx_bool &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Tensor &Tensor::operator/=<Scalar>(const Scalar &rc) {
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc, this->device()));
cytnx::linalg::iDiv(*this, _scalar_as_rank1_tensor(rc));
return *this;
}
template <>
Expand Down
33 changes: 30 additions & 3 deletions src/linalg/iAdd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,21 @@ namespace cytnx {
namespace linalg {

void iAdd(Tensor &Lt, const Tensor &Rt) {
cytnx_error_msg(Lt.device() != Rt.device(),
// A length-1 RHS that stays on the host is treated as a broadcast scalar: the GPU kernels
// read it with a host-side dereference and pass it into the kernel by value, so it needs
// neither a device match nor a per-call H2D copy of the scalar. See #988.
const bool rhs_is_host_scalar =
(Rt.device() == Device.cpu && Rt._impl->storage()._impl->size() == 1);
cytnx_error_msg(Lt.device() != Rt.device() && !rhs_is_host_scalar,
"[iAdd] error, the two tensors have to be on the same device.%s", "\n");
// In-place ops write the result back into the LHS storage, so a complex result cannot be
// stored in a real LHS. Guard here (device-independent) so the GPU path throws like the CPU
// path instead of reinterpreting the real output buffer as complex and corrupting it. See
// #988.
cytnx_error_msg(!Type.is_complex(Lt.dtype()) && Type.is_complex(Rt.dtype()),
"[iAdd] Cannot perform real += complex in-place: a complex result cannot be "
"stored in a real tensor.%s",
"\n");
if (!(Rt.shape().size() == 1 && Rt.shape()[0] == 1)) {
cytnx_error_msg(Lt.shape() != Rt.shape(),
"[iAdd] error, the two tensors do not have the same shape. Lt rank: [%d] "
Expand All @@ -27,6 +40,20 @@ namespace cytnx {
R = Rt;
}

// GPU broadcast scalar with a LHS *narrower* than the promoted dtype (e.g. a Float tensor
// plus a Double scalar, or an integer tensor plus a fractional scalar): the in-place GPU
// kernels write the promoted-width result straight into the narrower LHS buffer, an
// out-of-bounds write that corrupts memory. Compute in the promoted dtype, then truncate
// back to the LHS dtype -- this matches the CPU element-wise semantics (compute in the
// promoted type, store into the LHS type). See #988. The CPU path already does this
// in place, and real += complex is rejected above.
if (rhs_is_host_scalar && Lt.device() != Device.cpu && Lt.dtype() > Rt.dtype()) {
Tensor promoted = Lt.astype(Type.type_promote(Lt.dtype(), Rt.dtype()));
iAdd(promoted, R);
Lt = promoted.astype(Lt.dtype());
return;
}

static const std::vector<cytnx_uint64> empty_mapper;

// if contiguous, then no need to calculate the mappers
Expand All @@ -36,7 +63,7 @@ namespace cytnx {
detail::DispatchInplaceArithmeticCPU<0>(Lt, R, empty_mapper, empty_mapper, empty_mapper);
} else {
#ifdef UNI_GPU
checkCudaErrors(cudaSetDevice(Rt.device()));
checkCudaErrors(cudaSetDevice(Lt.device()));
Tensor tmpo;
if (Lt.dtype() <= Rt.dtype())
tmpo = Lt;
Expand All @@ -61,7 +88,7 @@ namespace cytnx {
Rt._impl->invmapper());
} else {
#ifdef UNI_GPU
checkCudaErrors(cudaSetDevice(Rt.device()));
checkCudaErrors(cudaSetDevice(Lt.device()));
Tensor tmpo;
if (Lt.dtype() <= Rt.dtype())
tmpo = Lt;
Expand Down
Loading
Loading