wehub-resource-sync
448 lines
14 KiB
C++
448 lines
14 KiB
C++
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <ATen/Functions.h>
|
|
#include <ATen/TensorIndexing.h>
|
|
#include <ATen/core/TensorBody.h>
|
|
#include <ATen/ops/tensor.h>
|
|
#include <c10/core/List.h>
|
|
#include <c10/core/ScalarType.h>
|
|
#include <c10/core/TensorOptions.h>
|
|
|
|
#include "ATen/ATen.h"
|
|
#include "gtest/gtest.h"
|
|
#include "paddle/common/macros.h"
|
|
#include "torch/all.h"
|
|
|
|
COMMON_DECLARE_bool(use_stride_kernel);
|
|
|
|
// ======================== index tests ========================
|
|
|
|
TEST(TensorIndexTest, IndexWithSingleTensor) {
|
|
// Create tensor [0, 10, 20, 30, 40]
|
|
at::Tensor t = at::arange(5, at::kFloat);
|
|
for (int i = 0; i < 5; i++) {
|
|
t.data_ptr<float>()[i] = static_cast<float>(i * 10);
|
|
}
|
|
|
|
// Index with [0, 2, 4]
|
|
at::Tensor idx = at::empty({3}, at::kLong);
|
|
int64_t* idx_data = idx.data_ptr<int64_t>();
|
|
idx_data[0] = 0;
|
|
idx_data[1] = 2;
|
|
idx_data[2] = 4;
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
at::Tensor result = t.index(indices);
|
|
ASSERT_EQ(result.numel(), 3);
|
|
|
|
float* result_data = result.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(result_data[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(result_data[1], 20.0f);
|
|
ASSERT_FLOAT_EQ(result_data[2], 40.0f);
|
|
}
|
|
|
|
TEST(TensorIndexTest, SliceKeepsStrideWithoutContiguousCopy) {
|
|
if (!FLAGS_use_stride_kernel) {
|
|
return;
|
|
}
|
|
at::Tensor base = at::arange(24, at::kFloat).reshape({4, 6});
|
|
at::Tensor transposed = base.t(); // shape: [6, 4], strides: [1, 6]
|
|
ASSERT_FALSE(transposed.is_contiguous());
|
|
|
|
at::Tensor sliced =
|
|
transposed.index({at::indexing::Slice(1, 5), at::indexing::Slice(0, 3)});
|
|
|
|
ASSERT_EQ(sliced.sizes(), c10::IntArrayRef({4, 3}));
|
|
ASSERT_EQ(sliced.strides(), c10::IntArrayRef({1, 6}));
|
|
ASSERT_EQ(sliced.stride(0), transposed.stride(0));
|
|
ASSERT_EQ(sliced.stride(1), transposed.stride(1));
|
|
ASSERT_FALSE(sliced.is_contiguous());
|
|
}
|
|
|
|
TEST(TensorIndexTest, IndexWithEmptyInitializerListReturnsSelf) {
|
|
at::Tensor t = at::arange(5, at::kFloat);
|
|
|
|
// PyTorch throws for empty index list
|
|
ASSERT_THROW(t.index(std::initializer_list<at::indexing::TensorIndex>{}),
|
|
std::exception);
|
|
}
|
|
|
|
TEST(TensorIndexTest, IndexWithTensorInitializerList) {
|
|
at::Tensor t = at::arange(5, at::kFloat);
|
|
|
|
at::Tensor idx = at::empty({3}, at::kLong);
|
|
int64_t* idx_data = idx.data_ptr<int64_t>();
|
|
idx_data[0] = 0;
|
|
idx_data[1] = 2;
|
|
idx_data[2] = 4;
|
|
|
|
at::Tensor result = t.index({idx});
|
|
|
|
ASSERT_EQ(result.numel(), 3);
|
|
float* result_data = result.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(result_data[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(result_data[1], 2.0f);
|
|
ASSERT_FLOAT_EQ(result_data[2], 4.0f);
|
|
}
|
|
|
|
TEST(TensorIndexTest, MemberIndexWithArrayRefTensorIndices) {
|
|
if (!FLAGS_use_stride_kernel) {
|
|
return;
|
|
}
|
|
at::Tensor base = at::arange(24, at::kFloat).reshape({4, 6});
|
|
at::Tensor transposed = base.t();
|
|
std::vector<at::indexing::TensorIndex> indices = {at::indexing::Slice(1, 5),
|
|
at::indexing::Slice(0, 3)};
|
|
|
|
at::Tensor sliced = transposed.index(indices);
|
|
|
|
ASSERT_EQ(sliced.sizes(), c10::IntArrayRef({4, 3}));
|
|
ASSERT_EQ(sliced.strides(), c10::IntArrayRef({1, 6}));
|
|
}
|
|
|
|
TEST(TensorIndexTest, MixedSliceAndTensorIndicesThrows) {
|
|
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 0;
|
|
idx.data_ptr<int64_t>()[1] = 2;
|
|
|
|
ASSERT_THROW(t.index({at::indexing::Slice(0, 2), idx}), std::exception);
|
|
}
|
|
|
|
// ======================== index_put_ tests ========================
|
|
|
|
TEST(TensorIndexPutTest, IndexPutInplaceWithTensor) {
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
float* original_data_ptr = t.data_ptr<float>();
|
|
|
|
// Create index tensor [1, 3]
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
int64_t* idx_data = idx.data_ptr<int64_t>();
|
|
idx_data[0] = 1;
|
|
idx_data[1] = 3;
|
|
|
|
// Values to put
|
|
at::Tensor values = at::full({2}, 99.0f, at::kFloat);
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
t.index_put_(indices, values);
|
|
|
|
// Verify data pointer unchanged (inplace)
|
|
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
|
|
|
float* data = t.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(data[1], 99.0f);
|
|
ASSERT_FLOAT_EQ(data[2], 0.0f);
|
|
ASSERT_FLOAT_EQ(data[3], 99.0f);
|
|
ASSERT_FLOAT_EQ(data[4], 0.0f);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutInplaceWithScalar) {
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
float* original_data_ptr = t.data_ptr<float>();
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
int64_t* idx_data = idx.data_ptr<int64_t>();
|
|
idx_data[0] = 0;
|
|
idx_data[1] = 4;
|
|
|
|
t.index_put_({idx}, at::Scalar(7.0));
|
|
|
|
// Verify data pointer unchanged (inplace)
|
|
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
|
|
|
float* data = t.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(data[0], 7.0f);
|
|
ASSERT_FLOAT_EQ(data[1], 0.0f);
|
|
ASSERT_FLOAT_EQ(data[4], 7.0f);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutNonInplace) {
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
int64_t* idx_data = idx.data_ptr<int64_t>();
|
|
idx_data[0] = 1;
|
|
idx_data[1] = 3;
|
|
|
|
at::Tensor values = at::full({2}, 42.0f, at::kFloat);
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
at::Tensor result = t.index_put(indices, values);
|
|
|
|
// Original should be unchanged
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 0.0f);
|
|
|
|
// Result should have the values
|
|
float* rdata = result.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(rdata[1], 42.0f);
|
|
ASSERT_FLOAT_EQ(rdata[3], 42.0f);
|
|
}
|
|
|
|
// ======================= Additional index edge case tests
|
|
// =======================
|
|
|
|
TEST(TensorIndexTest, IndexWithEmptyList) {
|
|
// Test index with empty indices list (should return self)
|
|
at::Tensor t = at::arange(5, at::kFloat);
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
|
|
at::Tensor result = t.index(indices);
|
|
ASSERT_EQ(result.numel(), 5);
|
|
}
|
|
|
|
TEST(TensorIndexTest, IndexWithMultipleIndices) {
|
|
// Test index with multiple indices (2D indexing)
|
|
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
|
|
|
at::Tensor idx0 = at::empty({2}, at::kLong);
|
|
int64_t* idx0_data = idx0.data_ptr<int64_t>();
|
|
idx0_data[0] = 0;
|
|
idx0_data[1] = 1;
|
|
|
|
at::Tensor idx1 = at::empty({2}, at::kLong);
|
|
int64_t* idx1_data = idx1.data_ptr<int64_t>();
|
|
idx1_data[0] = 0;
|
|
idx1_data[1] = 2;
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx0);
|
|
indices.push_back(idx1);
|
|
|
|
at::Tensor result = t.index(indices);
|
|
ASSERT_EQ(result.numel(), 2);
|
|
}
|
|
|
|
TEST(TensorIndexTest, IndexWithOptionalNone) {
|
|
// Test index with optional None in indices
|
|
// None means "select all" along that dimension
|
|
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 0;
|
|
idx.data_ptr<int64_t>()[1] = 2;
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(::std::nullopt); // None = select all rows
|
|
indices.push_back(idx); // [0, 2] = select columns 0 and 2
|
|
|
|
at::Tensor result = t.index(indices);
|
|
// Result should be shape {3, 2} = 6 elements
|
|
// Columns 0 and 2 from all rows: [[0,2], [3,5], [6,8]]
|
|
ASSERT_EQ(result.numel(), 6);
|
|
}
|
|
|
|
TEST(TensorIndexTest, FreeIndexWithAllNoneReturnsSelf) {
|
|
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(::std::nullopt);
|
|
indices.push_back(::std::nullopt);
|
|
|
|
at::Tensor result = at::index(t, indices);
|
|
|
|
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[5], 5.0f);
|
|
}
|
|
|
|
TEST(TensorIndexTest, FreeIndexWithSingleLeadingTensor) {
|
|
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 2;
|
|
idx.data_ptr<int64_t>()[1] = 0;
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
at::Tensor result = at::index(t, indices);
|
|
|
|
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 6.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 7.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[2], 8.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[3], 0.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[4], 1.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[5], 2.0f);
|
|
}
|
|
|
|
TEST(TensorIndexTest, MixedTensorNoneFullSliceIndex) {
|
|
at::Tensor base = at::arange(12, at::kFloat).reshape({3, 4});
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 2;
|
|
idx.data_ptr<int64_t>()[1] = 0;
|
|
|
|
at::Tensor result =
|
|
base.index({idx, at::indexing::None, at::indexing::Slice()});
|
|
|
|
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 1, 4}));
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 8.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 9.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[2], 10.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[3], 11.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[4], 0.0f);
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[7], 3.0f);
|
|
}
|
|
|
|
TEST(TensorIndexTest, MixedFullSliceWithMultipleTensorIndicesThrows) {
|
|
at::Tensor base = at::arange(12, at::kFloat).reshape({3, 4});
|
|
at::Tensor idx0 = at::empty({2}, at::kLong);
|
|
idx0.data_ptr<int64_t>()[0] = 0;
|
|
idx0.data_ptr<int64_t>()[1] = 1;
|
|
at::Tensor idx1 = at::empty({2}, at::kLong);
|
|
idx1.data_ptr<int64_t>()[0] = 0;
|
|
idx1.data_ptr<int64_t>()[1] = 1;
|
|
|
|
ASSERT_THROW(base.index({idx0, at::indexing::Slice(), idx1}), std::exception);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutAccumulate) {
|
|
// Test index_put_ with accumulate=true
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
float* original_data_ptr = t.data_ptr<float>();
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 1;
|
|
idx.data_ptr<int64_t>()[1] = 1;
|
|
|
|
at::Tensor values = at::full({2}, 5.0f, at::kFloat);
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
t.index_put_(indices, values, true); // accumulate=true
|
|
|
|
// Verify data pointer unchanged (inplace)
|
|
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
|
|
|
float* data = t.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(data[1], 10.0f); // 5 + 5 (accumulated)
|
|
ASSERT_FLOAT_EQ(data[2], 0.0f);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutWith2D) {
|
|
// Test index_put_ with 2D tensor
|
|
at::Tensor t = at::zeros({3, 3}, at::kFloat);
|
|
float* original_data_ptr = t.data_ptr<float>();
|
|
|
|
at::Tensor idx0 = at::arange(2, at::kLong);
|
|
idx0.data_ptr<int64_t>()[0] = 0;
|
|
idx0.data_ptr<int64_t>()[1] = 1;
|
|
at::Tensor idx1 = at::arange(2, at::kLong);
|
|
idx1.data_ptr<int64_t>()[0] = 0;
|
|
idx1.data_ptr<int64_t>()[1] = 1;
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx0);
|
|
indices.push_back(idx1);
|
|
|
|
at::Tensor values = at::full({2}, 9.0f, at::kFloat);
|
|
|
|
t.index_put_(indices, values);
|
|
|
|
// Verify data pointer unchanged (inplace)
|
|
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
|
|
|
float* data = t.data_ptr<float>();
|
|
ASSERT_FLOAT_EQ(data[0], 9.0f); // [0,0]
|
|
ASSERT_FLOAT_EQ(data[4], 9.0f); // [1,1]
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutNonInplaceAccumulate) {
|
|
// Test index_put with accumulate=true (non-inplace)
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 1;
|
|
idx.data_ptr<int64_t>()[1] = 1;
|
|
at::Tensor values = at::full({2}, 3.0f, at::kFloat);
|
|
|
|
c10::List<::std::optional<at::Tensor>> indices;
|
|
indices.push_back(idx);
|
|
|
|
at::Tensor result = t.index_put(indices, values, true);
|
|
|
|
// Original unchanged
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 0.0f);
|
|
// Result has accumulated
|
|
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 6.0f);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutArrayRefWithTensorValue) {
|
|
at::Tensor t = at::zeros({5}, at::kFloat);
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 1;
|
|
idx.data_ptr<int64_t>()[1] = 4;
|
|
at::Tensor values = at::full({2}, 13.0f, at::kFloat);
|
|
|
|
std::vector<at::indexing::TensorIndex> tensor_indices = {idx};
|
|
t.index_put_(at::ArrayRef<at::indexing::TensorIndex>(tensor_indices), values);
|
|
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[0], 0.0f);
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 13.0f);
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[2], 0.0f);
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[4], 13.0f);
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutArrayRefWithNoneValue) {
|
|
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
|
at::Tensor values = at::full({1, 2, 3}, 6.0f, at::kFloat);
|
|
|
|
t.index_put_({at::indexing::None}, values);
|
|
|
|
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 3}));
|
|
for (int i = 0; i < 6; ++i) {
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 6.0f);
|
|
}
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutArrayRefWithTensorNoneAndSlice) {
|
|
at::Tensor t = at::zeros({3, 4}, at::kFloat);
|
|
at::Tensor idx = at::empty({2}, at::kLong);
|
|
idx.data_ptr<int64_t>()[0] = 2;
|
|
idx.data_ptr<int64_t>()[1] = 0;
|
|
at::Tensor values = at::full({2, 1, 4}, 8.0f, at::kFloat);
|
|
|
|
t.index_put_({idx, at::indexing::None, at::indexing::Slice()}, values);
|
|
|
|
ASSERT_EQ(t.sizes(), c10::IntArrayRef({3, 4}));
|
|
for (int i = 0; i < 4; ++i) {
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 8.0f);
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[8 + i], 8.0f);
|
|
}
|
|
for (int i = 4; i < 8; ++i) {
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 0.0f);
|
|
}
|
|
}
|
|
|
|
TEST(TensorIndexPutTest, IndexPutArrayRefWithNoneScalarValue) {
|
|
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
|
|
|
t.index_put_({at::indexing::None}, at::Scalar(4.0));
|
|
|
|
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 3}));
|
|
for (int i = 0; i < 6; ++i) {
|
|
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 4.0f);
|
|
}
|
|
}
|