xref: /aosp_15_r20/external/pytorch/aten/src/ATen/test/memory_overlapping_test.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 
3 #include <ATen/ATen.h>
4 
5 using namespace at;
6 
7 std::vector<std::vector<int64_t>> sizes = {{1, 2, 3}, {1, 3, 2}, {2, 1, 3}, {3, 1, 2}, {3, 2, 1}, {2, 3, 1}};
8 
TEST(MemoryOverlapTest,TensorExpanded)9 TEST(MemoryOverlapTest, TensorExpanded) {
10   // NOLINTNEXTLINE(performance-for-range-copy)
11   for (auto size : sizes) {
12     Tensor t = at::ones({1}).expand(size);
13     EXPECT_FALSE(t.is_contiguous());
14     EXPECT_FALSE(t.is_non_overlapping_and_dense());
15   }
16 }
17 
TEST(MemoryOverlapTest,ScalarExpanded)18 TEST(MemoryOverlapTest, ScalarExpanded) {
19   // NOLINTNEXTLINE(performance-for-range-copy)
20   for (auto size : sizes) {
21     Tensor t = at::tensor(1).expand(size);
22     EXPECT_FALSE(t.is_contiguous());
23     EXPECT_FALSE(t.is_non_overlapping_and_dense());
24   }
25 }
26 
TEST(MemoryOverlapTest,NonContiguousTensor)27 TEST(MemoryOverlapTest, NonContiguousTensor) {
28   // NOLINTNEXTLINE(performance-for-range-copy)
29   for (auto size : sizes) {
30     Tensor t = at::rand(size).transpose(1, 2).transpose(0, 2);
31     if (!t.is_contiguous()) {
32       EXPECT_TRUE(t.is_non_overlapping_and_dense());
33     }
34   }
35 }
36 
TEST(MemoryOverlapTest,NonContiguousExpandedTensor)37 TEST(MemoryOverlapTest, NonContiguousExpandedTensor) {
38   // NOLINTNEXTLINE(performance-for-range-copy)
39   for (auto size : sizes) {
40     Tensor t = at::rand(size).transpose(1, 2).transpose(0, 2);
41     if (!t.is_contiguous()) {
42       for (auto size_to_add : {1, 2, 3, 4}) {
43         auto transpose_size = t.sizes().vec();
44         std::vector<int64_t> expanded_size(transpose_size);
45         expanded_size.insert(expanded_size.begin(), size_to_add);
46         auto expanded = t.expand(expanded_size);
47         EXPECT_FALSE(t.is_contiguous());
48         if (size_to_add == 1) {
49           EXPECT_TRUE(expanded.is_non_overlapping_and_dense());
50         } else {
51           EXPECT_FALSE(expanded.is_non_overlapping_and_dense());
52         }
53       }
54     }
55   }
56 }
57 
TEST(MemoryOverlapTest,ContiguousTensor)58 TEST(MemoryOverlapTest, ContiguousTensor) {
59   // NOLINTNEXTLINE(performance-for-range-copy)
60   for (auto size : sizes) {
61     Tensor t = at::rand(size);
62     EXPECT_TRUE(t.is_contiguous());
63     EXPECT_TRUE(t.is_non_overlapping_and_dense());
64   }
65 }
66 
TEST(MemoryOverlapTest,ContiguousExpandedTensor)67 TEST(MemoryOverlapTest, ContiguousExpandedTensor) {
68   // NOLINTNEXTLINE(performance-for-range-copy)
69   for (auto size : sizes) {
70     Tensor t = at::rand(size);
71     for (auto size_to_add : {1, 2, 3, 4}) {
72       std::vector<int64_t> expanded_size(size);
73       expanded_size.insert(expanded_size.begin(), size_to_add);
74       auto expanded = t.expand(expanded_size);
75       EXPECT_TRUE(t.is_contiguous());
76       EXPECT_TRUE(t.is_non_overlapping_and_dense());
77     }
78   }
79 }
80