[C++][Tensor] Incorrect logic for creating arrow::SparseCOOTensor from Tensor via arrow::SparseCOOTensor::Make (illegal memory access and incorrect values)

[andishgar]

Describe the bug, including details regarding any error messages, version, and platform.

The following method causes illegal memory access when creating from a combination of negative zero and non-zero values, and produces incorrect tensor values (potentially leading to illegal memory access) when creating from a column-major tensor.

arrow/cpp/src/arrow/sparse_tensor.h

Lines 583 to 594 in fddd356

	static inline Result<std::shared_ptr<SparseTensorImpl<SparseIndexType>>> Make(
	const Tensor& tensor, const std::shared_ptr<DataType>& index_value_type,
	MemoryPool* pool = default_memory_pool()) {
	std::shared_ptr<SparseIndex> sparse_index;
	std::shared_ptr<Buffer> data;
	ARROW_RETURN_NOT_OK(internal::MakeSparseTensorFromTensor(
	tensor, SparseIndexType::format_id, index_value_type, pool, &sparse_index,
	&data));
	return std::make_shared<SparseTensorImpl<SparseIndexType>>(
	internal::checked_pointer_cast<SparseIndexType>(sparse_index), tensor.type(),
	data, tensor.shape(), tensor.dim_names_);
	}

1- The following code leads to illegal memory access.

TEST(MyTest, SegFault) {
  // clang-format off
  std::vector<float> data{
    -0.0, -0.0, -0.0,
    -0.0, -0.0, -0.0,
    -0.0, -0.0, -0.0,
    -1.0, -0.0, -0.0,
    };

  // clang-format on

  std::vector<int64_t> shape = {4, 3};
  auto buffer = Buffer::FromVector(data);
  ASSERT_OK_AND_ASSIGN(auto dense_tensor, Tensor::Make(float32(), buffer, shape));
  ASSERT_OK_AND_ASSIGN(auto sparse_coo_tensor,
                       SparseCOOTensor::Make(*dense_tensor, int64()));
  ARROW_LOGGER_INFO("", sparse_coo_tensor->sparse_index()->non_zero_length());
}

and the error is:

: 1
mimalloc: error: buffer overflow in heap block 0x0200000100C0 of size 152: write after 152 bytes
Process finished with exit code 134 (interrupted by signal 6:SIGABRT)

The reason for this is rooted in the following code.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Line 179 in fddd356

ARROW_ASSIGN_OR_RAISE(int64_t nonzero_count, tensor_.CountNonZero());

In the above code, both positive and negative zero are considered zero; however, in the following code, negative zero is treated as a non-zero value.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Line 193 in fddd356

if (std::any_of(tensor_data, tensor_data + value_elsize, IsNonZero)) {

arrow/cpp/src/arrow/tensor/coo_converter.cc

Line 67 in fddd356

if (ARROW_PREDICT_FALSE(x != zero)) {

Note that in the above code, the type of c_value_type is always an unsigned integer due to the following code.

arrow/cpp/src/arrow/tensor/converter_internal.h

Lines 22 to 88 in fddd356

	#define DISPATCH(ACTION, index_elsize, value_elsize, ...) \
	switch (index_elsize) { \
	case 1: \
	switch (value_elsize) { \
	case 1: \
	ACTION(uint8_t, uint8_t, __VA_ARGS__); \
	break; \
	case 2: \
	ACTION(uint8_t, uint16_t, __VA_ARGS__); \
	break; \
	case 4: \
	ACTION(uint8_t, uint32_t, __VA_ARGS__); \
	break; \
	case 8: \
	ACTION(uint8_t, uint64_t, __VA_ARGS__); \
	break; \
	} \
	break; \
	case 2: \
	switch (value_elsize) { \
	case 1: \
	ACTION(uint16_t, uint8_t, __VA_ARGS__); \
	break; \
	case 2: \
	ACTION(uint16_t, uint16_t, __VA_ARGS__); \
	break; \
	case 4: \
	ACTION(uint16_t, uint32_t, __VA_ARGS__); \
	break; \
	case 8: \
	ACTION(uint16_t, uint64_t, __VA_ARGS__); \
	break; \
	} \
	break; \
	case 4: \
	switch (value_elsize) { \
	case 1: \
	ACTION(uint32_t, uint8_t, __VA_ARGS__); \
	break; \
	case 2: \
	ACTION(uint32_t, uint16_t, __VA_ARGS__); \
	break; \
	case 4: \
	ACTION(uint32_t, uint32_t, __VA_ARGS__); \
	break; \
	case 8: \
	ACTION(uint32_t, uint64_t, __VA_ARGS__); \
	break; \
	} \
	break; \
	case 8: \
	switch (value_elsize) { \
	case 1: \
	ACTION(int64_t, uint8_t, __VA_ARGS__); \
	break; \
	case 2: \
	ACTION(int64_t, uint16_t, __VA_ARGS__); \
	break; \
	case 4: \
	ACTION(int64_t, uint32_t, __VA_ARGS__); \
	break; \
	case 8: \
	ACTION(int64_t, uint64_t, __VA_ARGS__); \
	break; \
	} \
	break; \
	}

2- Incorrect tensor values when creating a SparseCOOTensor from a ColumnMajorTensor via SparseCOOTensor::Make

TEST(My, ColumnMajor) {
  // clang-format off
  std::vector<int> data{
    1, 4, 7, 10,
    2, 5, 8, 11,
    3, 6, 9, 12
                      };
  // clang-format on
  std::vector<int> data_2{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  auto buffer = Buffer::FromVector(data);
  auto buffer_2 = Buffer::FromVector(data_2);
  std::vector<int64_t> shape = {4, 3};
  std::vector<int64_t> strides = {sizeof(int), 4 * sizeof(int)};
  ASSERT_OK_AND_ASSIGN(auto tensor, Tensor::Make(int32(), buffer, shape, strides));
  ASSERT_OK_AND_ASSIGN(auto tensor_2, Tensor::Make(int32(), buffer_2, shape));
  ASSERT_TRUE(tensor->Equals(*tensor_2));
  ASSERT_TRUE(tensor->is_contiguous());
  ASSERT_TRUE(tensor->is_column_major());
  ASSERT_OK_AND_ASSIGN(auto sparse_tensor, SparseCOOTensor::Make(*tensor))
  ASSERT_OK_AND_ASSIGN(auto new_tensor, sparse_tensor->ToTensor());
  ASSERT_EQ(12, sparse_tensor->non_zero_length());
  ASSERT_TRUE(new_tensor->is_contiguous());
  ASSERT_TRUE(new_tensor->is_row_major());
  // new_tensor is not equal to tensor!!
  ASSERT_FALSE(new_tensor->Equals(*tensor));
}

ASSERT_FALSE(new_tensor->Equals(*tensor)); passes because of the following code.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Lines 86 to 91 in fddd356

	// transpose indices
	for (int64_t i = 0; i < size; ++i) {
	for (int j = 0; j < ndim / 2; ++j) {
	std::swap(indices[i * ndim + j], indices[i * ndim + ndim - j - 1]);
	}
	}

Given that logic, the index {2,3} is produced for a column-major contiguous tensor with shape {4,3}, which is not a correct index.

Component(s)

C++

[andishgar]

@pitrou, I would appreciate it if you could verify my observation.

[pitrou]

cc @rok

[andishgar]

Another error occurs when creating a SparseCOOTensor from a row-major tensor with a column-major SparseCOOIndex.

TEST(MyTest, SparseCOOTensorColumnMajorTest) {
  // clang-format off
  std::vector<int64_t> values = {1, 2, 3,
                                 4, 5, 6,
                                 7, 8, 9,
                                 10, 11, 12
    };
  std::vector<int64_t> CooIndices = {0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3,
                                     0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2};
  // clang-format on
  auto data_buffer = Buffer::FromVector(values);
  auto indices_buffer = Buffer::FromVector(CooIndices);
  ASSERT_OK_AND_ASSIGN(auto indices_tensor,
                       Tensor::Make(int64(), indices_buffer, {12, 2}, {8, 96}));
  ASSERT_TRUE(indices_tensor->is_column_major());
  ASSERT_OK_AND_ASSIGN(auto sparse_coo_index, SparseCOOIndex::Make(indices_tensor, true));
  auto data_pointer = reinterpret_cast<const int64_t*>(data_buffer->data());
  auto tensor = Tensor::Make(int64(), data_buffer, {4, 3}).ValueOrDie();
  ASSERT_TRUE(tensor->is_row_major());
  SparseCOOTensor sparse_tensor(sparse_coo_index, int64(), data_buffer, {4,3}, {});
  // SparseCooIndex is stored as column major tensor
  ASSERT_TRUE(internal::checked_pointer_cast<SparseCOOIndex>(sparse_tensor.sparse_index())
                  ->indices()
                  ->is_column_major());
  // Check the correctness of indices
  for (int64_t i = 0; i < 12; i++) {
    auto row =
        internal::checked_pointer_cast<SparseCOOIndex>(sparse_tensor.sparse_index())
            ->indices()
            ->Value<Int64Type>({i, 0});
    auto column =
        internal::checked_pointer_cast<SparseCOOIndex>(sparse_tensor.sparse_index())
            ->indices()
            ->Value<Int64Type>({i, 1});
    ASSERT_TRUE(data_pointer[i] == tensor->Value<Int64Type>({row, column}));
  }
  auto new_tensor = sparse_tensor.ToTensor().ValueOrDie();
  ASSERT_TRUE(new_tensor->is_row_major());
  ASSERT_FALSE(new_tensor->Equals(*tensor));

ASSERT_FALSE(new_tensor->Equals(*tensor)); passed because it is based on the implementation below, which treats SparseCOOIndex as a row-major tensor.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Lines 316 to 321 in 982d31f

	for (int j = 0; j < ndim; ++j) {
	auto index = static_cast<int64_t>(
	SparseTensorConverterMixin::GetIndexValue(coords_data, index_elsize));
	offset += index * strides[j];
	coords_data += index_elsize;
	}

Note that, according to the following test, it is acceptable to have a column-major SparseCOOIndex.

arrow/cpp/src/arrow/sparse_tensor_test.cc

Lines 109 to 123 in 982d31f

	TEST(TestSparseCOOIndex, MakeColumnMajorCanonical) {
	std::vector<int32_t> values = {0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 2, 2,
	0, 0, 1, 1, 2, 2, 0, 2, 1, 3, 0, 2, 1, 3, 0, 2, 1, 3};
	auto data = Buffer::Wrap(values);
	std::vector<int64_t> shape = {12, 3};
	std::vector<int64_t> strides = {sizeof(int32_t), 12 * sizeof(int32_t)}; // Column-major
	// OK
	std::shared_ptr<SparseCOOIndex> si;
	ASSERT_OK_AND_ASSIGN(si, SparseCOOIndex::Make(int32(), shape, strides, data));
	ASSERT_EQ(shape, si->indices()->shape());
	ASSERT_EQ(strides, si->indices()->strides());
	ASSERT_EQ(data->data(), si->indices()->raw_data());
	ASSERT_TRUE(si->is_canonical());
	}

[andishgar]

@rok Is it okay to change the behavior of the following method to return a column-major tensor whenever necessary? (This refers to the second case mentioned in this comment.)

arrow/cpp/src/arrow/sparse_tensor.h

Lines 503 to 509 in 982d31f

	/// \brief Return dense representation of sparse tensor as tensor
	///
	/// The returned Tensor has row-major order (C-like).
	Result<std::shared_ptr<Tensor>> ToTensor(MemoryPool* pool) const;
	Result<std::shared_ptr<Tensor>> ToTensor() const {
	return ToTensor(default_memory_pool());
	}

[rok]

@andishgar thanks for finding and pointing these out!

In the above code, both positive and negative zero are considered zero; however, in the following code, negative zero is treated as a non-zero value.

This is quite unfortunate and we should address it asap. @pitrou it would probably be best to fix this before the release.

@rok Is it okay to change the behavior of the following method to return a column-major tensor whenever necessary? (This refers to the second case mentioned in this #47520 (comment).)

I think that'd be ok, what kind of API would you propose?
I imagine you wouldn't be satisfied with Tensor.to_column_major() that would make a data copy?

[andishgar]

I think the current API for converting SparseCooTensor to Tensor is okay; however, the implementation forgot to address two things:

1-SparseCooIndex can be stored in either column-major or row-major order, and this can be determined from SparseCooIndex itself since it is essentially a tensor. However, the current implementation assumes that SparseCooIndex is always row-major.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Lines 316 to 321 in 982d31f

	for (int j = 0; j < ndim; ++j) {
	auto index = static_cast<int64_t>(
	SparseTensorConverterMixin::GetIndexValue(coords_data, index_elsize));
	offset += index * strides[j];
	coords_data += index_elsize;
	}

2-These coordinates should be multiplied by the appropriate column- or row-major strides to determine the data locations in the final tensor. We also need to maintain a variable to track this stride and apply it to the result tensor. However, the current implementation assumes that the final stride is always row-major.

arrow/cpp/src/arrow/tensor/coo_converter.cc

Line 308 in 982d31f

RETURN_NOT_OK(ComputeRowMajorStrides(value_type, sparse_tensor->shape(), &strides));

arrow/cpp/src/arrow/tensor/coo_converter.cc

Line 319 in 982d31f

offset += index * strides[j];

arrow/cpp/src/arrow/tensor/coo_converter.cc

Lines 327 to 329 in 982d31f

	return std::make_shared<Tensor>(sparse_tensor->type(), std::move(values_buffer),
	sparse_tensor->shape(), strides,
	sparse_tensor->dim_names());

Let me investigate other sparse matrix formats (CSC, CSR, CSF) to see if they have the same logical error as case 2. Then I’ll give my final thoughts.

[andishgar]

@rok I've investigated all tensor sparse formats and I want to give my final opinion:

1. Regarding this, my assumption about case 2 was wrong; the problem only exists in the CooTensor format.
2. All types of sparse tensor format creation from dense tensors have the problem of handling negative zero.
3. HalfFloatType is not handled during the creation of any format at all (negative zero does not cause a memory problem).

This is quite unfortunate, and we should address it ASAP. pitrou, it would probably be best to fix this before the release.

My initial solution is ready; however, since I've changed 12 files, it will take several days to clean it up. I think the release is frozen on October 6th.

[andishgar]

@rok @pitrou
One more point regarding the sparse tensor specification: the current format treats all negative zeros as positive zeros. So, converting a dense tensor to a sparse tensor and back will turn every negative zero into a positive zero. Is this intended behavior?

[rok]

@rok I've investigated all tensor sparse formats and I want to give my final opinion:

Thanks for checking!

1. Regarding this, my assumption about case 2 was wrong; the problem only exists in the CooTensor format.

Great to hear it's a COO only issue.

2. All types of sparse tensor format creation from dense tensors have the problem of handling negative zero.

This is not great, but not sure how and if we should be fixing it. It's definitely is less urgent than the COO issue. Let's open a new issue for it?

3. HalfFloatType is not handled during the creation of any format at all (negative zero does not cause a memory problem).

I'd say this is a new feature and we'd want to handle it separately in a new issue.

This is quite unfortunate, and we should address it ASAP. pitrou, it would probably be best to fix this before the release.

My initial solution is ready; however, since I've changed 12 files, it will take several days to clean it up. I think the release is frozen on October 6th.

Do you have an open PR?

@rok @pitrou One more point regarding the sparse tensor specification: the current format treats all negative zeros as positive zeros. So, converting a dense tensor to a sparse tensor and back will turn every negative zero into a positive zero. Is this intended behavior?

This is the same as 2. ? It's not intended behavior if I remember correctly. I think this didn't come up when we originally implemented these. Perhaps we should document it.

[rok]

@andishgar if you don't have time for the COO issue before the release I can maybe take a look at it.

[andishgar]

@rok

This is not great, but not sure how and if we should be fixing it. It's definitely is less urgent than the COO issue. Let's open a new issue for it?
I'd say this is a new feature and we'd want to handle it separately in a new issue.
Do you have an open PR?

My solution is ready. I have addressed negative zero, HalfFloat, and column-major tensors in all sparse tensor formats. I just need to clean up and write a test for dense tensor HalfFloat comparison. Should I spilit it to several new issues?

[rok]

@andishgar your patch looks good but it does cover a lot of ground. It would be nice to have wrong-index-when-positive/negative-zero, column-major and half-float features as separate PRs. Perhaps you could start with opening a PR for this issue adding just is_not_zero utility (without float16) so we can land it before the release? It seems like it should be a small change.
For the rest of the diff perhaps first make an issue for the column-major index feature and use the full patch. Float16 would also be nice in a separate patch, but it'll probably be easier to split out than the column-major part.