Files
2026-07-13 12:40:42 +08:00

209 lines
7.0 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/core/TensorBody.h>
#include <ATen/cuda/EmptyTensor.h>
#include <ATen/native/cuda/Resize.h>
#include <ATen/ops/tensor.h>
#include <c10/core/ScalarType.h>
#include <c10/core/SymInt.h>
#include <c10/core/TensorOptions.h>
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <c10/cuda/CUDAFunctions.h>
#include <c10/cuda/CUDAGuard.h>
#endif
#include "ATen/ATen.h"
#include "gtest/gtest.h"
#include "paddle/phi/common/float16.h"
#include "torch/all.h"
// ============================================================================
// Narrow Tests
// ============================================================================
TEST(TestNarrow, NarrowBasic) {
// Test basic narrow operation on dimension 0
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
// tensor: [[0, 1, 2, 3],
// [4, 5, 6, 7],
// [8, 9, 10, 11]]
// Narrow on dim 0, start 1, length 2
at::Tensor narrowed = tensor.narrow(0, 1, 2);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 4}));
// Verify data
const float* data = narrowed.data_ptr<float>();
ASSERT_EQ(data[0], 4.0f); // First element of row 1
ASSERT_EQ(data[4], 8.0f); // First element of row 2
}
TEST(TestNarrow, NarrowOnDim1) {
// Test narrow on dimension 1
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
// Narrow on dim 1, start 1, length 2
at::Tensor narrowed = tensor.narrow(1, 1, 2);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({3, 2}));
}
TEST(TestNarrow, NarrowStartZero) {
// Test narrow starting from 0
at::Tensor tensor = at::ones({5, 4}, at::kFloat);
at::Tensor narrowed = tensor.narrow(0, 0, 3);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({3, 4}));
ASSERT_EQ(narrowed.numel(), 12);
}
TEST(TestNarrow, NarrowFullLength) {
// Test narrow with full length (should return same shape on that dim)
at::Tensor tensor = at::ones({3, 4}, at::kFloat);
at::Tensor narrowed = tensor.narrow(0, 0, 3);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({3, 4}));
}
TEST(TestNarrow, NarrowSingleElement) {
// Test narrow with length 1
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
at::Tensor narrowed = tensor.narrow(0, 1, 1);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({1, 4}));
}
TEST(TestNarrow, NarrowHigherDim) {
// Test narrow on 3D tensor
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
// Narrow on dim 1
at::Tensor narrowed = tensor.narrow(1, 1, 2);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 2, 4}));
// Narrow on dim 2
at::Tensor narrowed2 = tensor.narrow(2, 0, 2);
ASSERT_EQ(narrowed2.sizes(), c10::IntArrayRef({2, 3, 2}));
}
TEST(TestNarrow, NarrowSymInt) {
// Test narrow_symint (should behave same as narrow)
at::Tensor tensor = at::ones({3, 4}, at::kFloat);
c10::SymInt start = 1;
c10::SymInt length = 2;
at::Tensor narrowed = tensor.narrow_symint(0, start, length);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 4}));
}
// ============================================================================
// Narrow Copy Tests
// ============================================================================
TEST(TestNarrowCopy, NarrowCopyBasic) {
// Test narrow_copy returns a copy of the narrowed tensor
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
at::Tensor narrowed_copy = tensor.narrow_copy(0, 1, 2);
ASSERT_EQ(narrowed_copy.sizes(), c10::IntArrayRef({2, 4}));
// Verify it's a copy (different data pointer)
// Note: narrow_copy creates a contiguous copy
ASSERT_EQ(narrowed_copy.numel(), 8);
}
TEST(TestNarrowCopy, NarrowCopySymInt) {
// Test narrow_copy_symint
at::Tensor tensor = at::ones({3, 4}, at::kFloat);
c10::SymInt start = 0;
c10::SymInt length = 2;
at::Tensor narrowed_copy = tensor.narrow_copy_symint(0, start, length);
ASSERT_EQ(narrowed_copy.sizes(), c10::IntArrayRef({2, 4}));
}
TEST(TestNarrowCopy, NarrowCopyDataIntegrity) {
// Test that narrow_copy preserves data correctly
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
at::Tensor narrowed_copy = tensor.narrow_copy(0, 1, 2);
// Verify data
const float* data = narrowed_copy.data_ptr<float>();
ASSERT_EQ(data[0], 4.0f);
ASSERT_EQ(data[1], 5.0f);
ASSERT_EQ(data[2], 6.0f);
ASSERT_EQ(data[3], 7.0f);
}
// ============================================================================
// Narrow with Tensor Start Tests
// ============================================================================
TEST(TestNarrowTensor, NarrowWithTensorStart) {
// Test narrow with tensor start parameter
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
// Create start tensor with value 1 using ones
at::Tensor start_tensor = at::ones({1}, at::kLong);
at::Tensor narrowed = tensor.narrow(0, start_tensor, 2);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 4}));
}
TEST(TestNarrowTensor, NarrowSymIntWithTensorStart) {
// Test narrow_symint with tensor start parameter
at::Tensor tensor = at::ones({3, 4}, at::kFloat);
// Create start tensor with value 0 using zeros
at::Tensor start_tensor = at::zeros({1}, at::kLong);
c10::SymInt length = 2;
at::Tensor narrowed = tensor.narrow_symint(0, start_tensor, length);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 4}));
}
// ============================================================================
// Combined Tests
// ============================================================================
TEST(TestNarrowCombined, NarrowMultipleDims) {
// Test narrow on multiple dimensions sequentially
at::Tensor tensor = at::ones({4, 5, 6}, at::kFloat);
// Narrow on each dimension
at::Tensor n1 = tensor.narrow(0, 1, 2);
ASSERT_EQ(n1.sizes(), c10::IntArrayRef({2, 5, 6}));
at::Tensor n2 = n1.narrow(1, 2, 3);
ASSERT_EQ(n2.sizes(), c10::IntArrayRef({2, 3, 6}));
at::Tensor n3 = n2.narrow(2, 0, 4);
ASSERT_EQ(n3.sizes(), c10::IntArrayRef({2, 3, 4}));
}
TEST(TestNarrowCombined, NarrowDataVerification) {
// Verify narrow selects correct data
at::Tensor tensor = at::arange(20, at::kFloat).reshape({4, 5});
// tensor: [[ 0, 1, 2, 3, 4],
// [ 5, 6, 7, 8, 9],
// [10, 11, 12, 13, 14],
// [15, 16, 17, 18, 19]]
// Narrow rows 1-2, columns 2-3
at::Tensor narrowed = tensor.narrow(0, 1, 2).narrow(1, 2, 2);
ASSERT_EQ(narrowed.sizes(), c10::IntArrayRef({2, 2}));
// Expected: [[7, 8], [12, 13]]
// Note: The actual verification depends on contiguous memory layout
}