Skip to content

fix: prevent integer overflow in Psi offsets #8107

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
mohanchen merged 1 commit into from
Oct 10, 2026
Merged

fix: prevent integer overflow in Psi offsets #8107

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
27 changes: 15 additions & 12 deletions source/source_psi/psi.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ T* Psi<T, Device>::get_pointer(const int& ikb) const
{
assert(ikb >= 0);
assert(this->k_first ? ikb < this->nbands : ikb < this->nk);
return this->psi_current + ikb * this->nbasis;
return this->psi_current + static_cast<std::size_t>(ikb) * this->nbasis;
}

template <typename T, typename Device>
Expand Down Expand Up @@ -377,7 +377,7 @@ void Psi<T, Device>::fix_k(const int ik) const
{
this->current_b = 0;
}
int base = this->current_b * this->nk * this->nbasis;
const std::size_t base = static_cast<std::size_t>(this->current_b) * this->nk * this->nbasis;
if (ik >= this->nk)
{
// mem_saver: fix to base
Expand All @@ -386,7 +386,8 @@ void Psi<T, Device>::fix_k(const int ik) const
}
else
{
this->psi_bias = k_first ? ik * this->nbands * this->nbasis : base + ik * this->nbasis;
this->psi_bias = k_first ? static_cast<std::size_t>(ik) * this->nbands * this->nbasis
: base + static_cast<std::size_t>(ik) * this->nbasis;
this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
}
}
Expand All @@ -400,7 +401,7 @@ void Psi<T, Device>::fix_b(const int ib) const
{
this->current_k = 0;
}
int base = this->current_k * this->nbands * this->nbasis;
const std::size_t base = static_cast<std::size_t>(this->current_k) * this->nbands * this->nbasis;
if (ib >= this->nbands)
{
// mem_saver: fix to base
Expand All @@ -409,7 +410,8 @@ void Psi<T, Device>::fix_b(const int ib) const
}
else
{
this->psi_bias = k_first ? base + ib * this->nbasis : ib * this->nk * this->nbasis;
this->psi_bias = k_first ? base + static_cast<std::size_t>(ib) * this->nbasis
: static_cast<std::size_t>(ib) * this->nk * this->nbasis;
this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
}
}
Expand All @@ -427,7 +429,8 @@ void Psi<T, Device>::fix_kb(const int ik, const int ib) const
}
else
{
this->psi_bias = k_first ? (ik * this->nbands + ib) * this->nbasis : (ib * this->nk + ik) * this->nbasis;
this->psi_bias = k_first ? (static_cast<std::size_t>(ik) * this->nbands + ib) * this->nbasis
: (static_cast<std::size_t>(ib) * this->nk + ik) * this->nbasis;
this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
}
}
Expand All @@ -437,8 +440,8 @@ T& Psi<T, Device>::operator()(const int ikb1, const int ikb2, const int ibasis)
{
assert(ikb1 >= 0 && ikb2 >= 0 && ibasis >= 0);
assert(this->k_first ? ikb1 < this->nk && ikb2 < this->nbands : ikb1 < this->nbands && ikb2 < this->nk);
return this->k_first ? this->psi[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis]
: this->psi[(ikb1 * this->nk + ikb2) * this->nbasis + ibasis];
return this->k_first ? this->psi[(static_cast<std::size_t>(ikb1) * this->nbands + ikb2) * this->nbasis + ibasis]
: this->psi[(static_cast<std::size_t>(ikb1) * this->nk + ikb2) * this->nbasis + ibasis];
}

template <typename T, typename Device>
Expand All @@ -447,7 +450,7 @@ T& Psi<T, Device>::operator()(const int ikb2, const int ibasis) const
assert(this->k_first ? this->current_b == 0 : this->current_k == 0);
assert(this->k_first ? ikb2 >= 0 && ikb2 < this->nbands : ikb2 >= 0 && ikb2 < this->nk);
assert(ibasis >= 0 && ibasis < this->nbasis);
return this->psi_current[ikb2 * this->nbasis + ibasis];
return this->psi_current[static_cast<std::size_t>(ikb2) * this->nbasis + ibasis];
}

template <typename T, typename Device>
Expand Down Expand Up @@ -508,13 +511,13 @@ std::tuple<const T*, int> Psi<T, Device>::to_range(const Range& range) const
}
else if (i1 < 0) // [r1, r2] is the range of index1 with length m
{
const T* p = &this->psi[r1 * (k_first ? this->nbands : this->nk) * this->nbasis];
const T* p = &this->psi[static_cast<std::size_t>(r1) * (k_first ? this->nbands : this->nk) * this->nbasis];
int m = (r2 - r1 + 1) * this->get_npol();
return std::tuple<const T*, int>(p, m);
}
else // [r1, r2] is the range of index2 with length m
{
const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
const T* p = &this->psi[(static_cast<std::size_t>(i1) * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
int m = (r2 - r1 + 1) * this->get_npol();
return std::tuple<const T*, int>(p, m);
}
Expand Down Expand Up @@ -553,4 +556,4 @@ template Psi<std::complex<float>, base_device::DEVICE_GPU>::Psi(
template Psi<std::complex<double>, base_device::DEVICE_GPU>::Psi(
const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
#endif
} // namespace psi
} // namespace psi
69 changes: 69 additions & 0 deletions source/source_psi/test/psi_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,75 @@
#include "source_psi/psi.h"

#include <gtest/gtest.h>
#include <memory>

TEST(PsiLargeOffsets, ThreeDimensionalAccess)
{
// No zero-initialization: the test touches only one page of this 12 GiB buffer.
std::unique_ptr<float[]> storage(new float[3221225472ULL]);
float* const data = storage.get();
data[2684354567ULL] = 42.0f;
for (const bool k_first : {true, false})
{
const int nk = k_first ? 3 : 2;
const int nbands = k_first ? 2 : 3;
const int ik = k_first ? 2 : 1;
const int ib = k_first ? 1 : 2;
psi::Psi<float> wave(data, nk, nbands, 536870912, 536870912, k_first);

if (k_first)
{
wave.fix_k(ik);
}
else
{
wave.fix_b(ib);
}
ASSERT_EQ(wave.get_psi_bias(), 2147483648ULL);
ASSERT_EQ(wave.get_pointer(), data + 2147483648ULL);
ASSERT_EQ(wave.get_pointer(1), data + 2684354560ULL);

if (k_first)
{
wave.fix_b(ib);
}
else
{
wave.fix_k(ik);
}
EXPECT_EQ(wave.get_pointer(), data + 2684354560ULL);

wave.fix_kb(ik, ib);
ASSERT_EQ(wave.get_psi_bias(), 2684354560ULL);
EXPECT_EQ(wave.get_pointer(), data + 2684354560ULL);
ASSERT_EQ(&wave(2, 1, 7), data + 2684354567ULL);
EXPECT_FLOAT_EQ(wave(2, 1, 7), 42.0f);
EXPECT_EQ(&wave(1, 1, 536870911), data + 2147483647ULL);

const psi::Range range(k_first, 2, 1, 1);
const auto result = wave.to_range(range);
ASSERT_EQ(std::get<0>(result), data + 2684354560ULL);
const psi::Range outer_range(k_first, -1, 2, 2);
const auto outer_result = wave.to_range(outer_range);
EXPECT_EQ(std::get<0>(outer_result), data + 2147483648ULL);
}
}

TEST(PsiLargeOffsets, TwoDimensionalAccess)
{
std::unique_ptr<float[]> storage(new float[3221225472ULL]);
float* const data = storage.get();
data[2147483655ULL] = 24.0f;
for (const bool k_first : {true, false})
{
const int nk = k_first ? 1 : 3;
const int nbands = k_first ? 3 : 1;
psi::Psi<float> wave(data, nk, nbands, 1073741824, 1073741824, k_first);
ASSERT_EQ(wave.get_pointer(2), data + 2147483648ULL);
ASSERT_EQ(&wave(2, 7), data + 2147483655ULL);
EXPECT_FLOAT_EQ(wave(2, 7), 24.0f);
}
}

class TestPsi : public ::testing::Test
{
Expand Down
Loading