Skip to content
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

Jump to bottom

Incorrect calculation of num_triplets at SparseMatrix construction #8501

Open
junwha opened this issue Mar 31, 2024 · 0 comments · May be fixed by #8502
Open

Incorrect calculation of num_triplets at SparseMatrix construction #8501

junwha opened this issue Mar 31, 2024 · 0 comments · May be fixed by #8502

Comments

@junwha
Copy link

junwha commented Mar 31, 2024 •
edited
Loading

Describe the bug
Sparse matrix is constructed with incorrect num_triplets, which leads to buffer-overflow read.

To Reproduce
PoC was modified from test_sparse_matrix.py

import taichi as ti
arch = ti.cpu # or ti.cuda
ti.init(arch=arch)

def test_build_sparse_matrix_frome_ndarray(dtype, storage_format):
    n = 8
    triplets = ti.Vector.ndarray(n=3, dtype=ti.f32, shape=n)
    A = ti.linalg.SparseMatrix(n=10, m=10, dtype=ti.f32, storage_format=storage_format)

    @ti.kernel
    def fill(triplets: ti.types.ndarray()):
        for i in range(n):
            triplet = ti.Vector([i, i, i], dt=ti.f32)
            triplets[i] = triplet

    fill(triplets)
    A.build_from_ndarray(triplets)

    for i in range(n):
        assert A[i, i] == i

test_build_sparse_matrix_frome_ndarray(ti.f32, "col_major")

Additional comments
At make_sparse_matrix_from_ndarray (taichi/program/sparse_matrix.cpp:378), it calculates num_triplets by ndarray.get_nelement() * ndarray.get_element_size() / 3. Here, let ndarray.get_nelement() be N and ndarray.get_element_size() be M. and we know only 3*N*M bytes are accessible from data_ptr.

void make_sparse_matrix_from_ndarray(Program *prog,
                                     SparseMatrix &sm,
                                     const Ndarray &ndarray) {
  std::string sdtype = taichi::lang::data_type_name(sm.get_data_type());
  auto data_ptr = prog->get_ndarray_data_ptr_as_int(&ndarray);
  auto num_triplets = ndarray.get_nelement() * ndarray.get_element_size() / 3;
  if (sdtype == "f32") {
    build_ndarray_template<float32>(sm, data_ptr, num_triplets);
  } else if (sdtype == "f64") {
    build_ndarray_template<float64>(sm, data_ptr, num_triplets);
  } else {
    TI_ERROR("Unsupported sparse matrix data type {}!", sdtype);
  }
}

And at build_ndarray_template (taichi/program/sparse_matrix.cpp:373), it casts data to T typed array, and accesses to index 0 to 3*(num_triplets-1)+2, which is 3*(N*M/3-1)+2 = N*M-1.
Thus, it accesses ((char*) data_ptr + (N*M-1)*M), that is, it overflows the limit ((T*) data_ptr + 3*N-1).

template <typename T>
void build_ndarray_template(SparseMatrix &sm,
                            intptr_t data_ptr,
                            size_t num_triplets) {
  using V = Eigen::Triplet<T>;
  std::vector<V> triplets;
  T *data = reinterpret_cast<T *>(data_ptr);
  for (int i = 0; i < num_triplets; i++) {
    x.push_back(
        V(data[i * 3], data[i * 3 + 1], taichi_union_cast<T>(data[i * 3 + 2])));
  }
  sm.build_triplets(static_cast<void *>(&triplets));
}

To fix this, we need to correct the num_triplets as ndarray.get_nelement().
I will open the PR for this.

Thank you!:)
@junwha junwha linked a pull request Mar 31, 2024 that will close this issue
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
Taichi Lang
Status: Untriaged
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

[Bug] Use correct num_triplets junwha/taichi
1 participant
@junwha