Skip to content

Commit 59778bf

Browse files
pytorchbotGasoonjia
pytorchbot
and
Gasoonjia
authored
[slimtensor] Add SizesAndStrides and dimension utility functions for slimtensor usage (#16549)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #16384 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/gasoonjia/75/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/75/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/gasoonjia/74/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/75/orig Differential Revision: [D89749336](https://our.internmc.facebook.com/intern/diff/D89749336/) @diff-train-skip-merge --------- Co-authored-by: gasoonjia <[email protected]> Co-authored-by: Gasoonjia <[email protected]>
1 parent 3c10c89 commit 59778bf

File tree

13 files changed

+1205
-0
lines changed

13 files changed

+1205
-0
lines changed

‎backends/aoti/slim/c10/core/SizesAndStrides.h‎

Lines changed: 415 additions & 0 deletions
Large diffs are not rendered by default.

‎backends/aoti/slim/c10/core/targets.bzl‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,20 @@ def define_common_targets():
4040
],
4141
)
4242

43+
# Header-only library for SizesAndStrides
44+
runtime.cxx_library(
45+
name = "sizes_and_strides",
46+
headers = [
47+
"SizesAndStrides.h",
48+
],
49+
visibility = ["@EXECUTORCH_CLIENTS"],
50+
exported_deps = [
51+
"//executorch/backends/aoti/slim/c10/macros:macros",
52+
"//executorch/runtime/core:core",
53+
"//executorch/runtime/platform:platform",
54+
],
55+
)
56+
4357
# Combined c10 core library
4458
runtime.cxx_library(
4559
name = "core",
@@ -48,5 +62,6 @@ def define_common_targets():
4862
":device",
4963
":device_type",
5064
":scalar_type",
65+
":sizes_and_strides",
5166
],
5267
)

‎backends/aoti/slim/c10/core/test/targets.bzl‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,13 @@ def define_common_targets():
2323
"//executorch/backends/aoti/slim/c10/core:scalar_type",
2424
],
2525
)
26+
27+
runtime.cxx_test(
28+
name = "test_sizes_and_strides",
29+
srcs = [
30+
"test_sizes_and_strides.cpp",
31+
],
32+
deps = [
33+
"//executorch/backends/aoti/slim/c10/core:sizes_and_strides",
34+
],
35+
)

‎backends/aoti/slim/c10/core/test/test_sizes_and_strides.cpp‎

Lines changed: 316 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,316 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <gtest/gtest.h>
10+
11+
#include <executorch/backends/aoti/slim/c10/core/SizesAndStrides.h>
12+
13+
namespace executorch::backends::aoti::slim::c10 {
14+
15+
// =============================================================================
16+
// Default Construction Tests
17+
// =============================================================================
18+
19+
TEST(SizesAndStridesTest, DefaultConstruct) {
20+
SizesAndStrides ss;
21+
22+
EXPECT_EQ(ss.size(), 1u);
23+
EXPECT_EQ(ss.size_at(0), 0);
24+
EXPECT_EQ(ss.stride_at(0), 1);
25+
}
26+
27+
// =============================================================================
28+
// Set Sizes and Strides Tests
29+
// =============================================================================
30+
31+
TEST(SizesAndStridesTest, SetSizes) {
32+
SizesAndStrides ss;
33+
ss.set_sizes({2, 3, 4});
34+
35+
EXPECT_EQ(ss.size(), 3u);
36+
EXPECT_EQ(ss.sizes_arrayref()[0], 2);
37+
EXPECT_EQ(ss.sizes_arrayref()[1], 3);
38+
EXPECT_EQ(ss.sizes_arrayref()[2], 4);
39+
}
40+
41+
TEST(SizesAndStridesTest, SetSizesAndStrides) {
42+
SizesAndStrides ss;
43+
ss.set_sizes({2, 3, 4});
44+
ss.set_strides({12, 4, 1});
45+
46+
EXPECT_EQ(ss.size(), 3u);
47+
EXPECT_EQ(ss.sizes_arrayref()[0], 2);
48+
EXPECT_EQ(ss.sizes_arrayref()[1], 3);
49+
EXPECT_EQ(ss.sizes_arrayref()[2], 4);
50+
EXPECT_EQ(ss.strides_arrayref()[0], 12);
51+
EXPECT_EQ(ss.strides_arrayref()[1], 4);
52+
EXPECT_EQ(ss.strides_arrayref()[2], 1);
53+
}
54+
55+
TEST(SizesAndStridesTest, SizeAt) {
56+
SizesAndStrides ss;
57+
ss.set_sizes({10, 20, 30});
58+
59+
EXPECT_EQ(ss.size_at(0), 10);
60+
EXPECT_EQ(ss.size_at(1), 20);
61+
EXPECT_EQ(ss.size_at(2), 30);
62+
}
63+
64+
TEST(SizesAndStridesTest, StrideAt) {
65+
SizesAndStrides ss;
66+
ss.set_sizes({2, 3});
67+
ss.set_strides({3, 1});
68+
69+
EXPECT_EQ(ss.stride_at(0), 3);
70+
EXPECT_EQ(ss.stride_at(1), 1);
71+
}
72+
73+
TEST(SizesAndStridesTest, SizeAtUnchecked) {
74+
SizesAndStrides ss;
75+
ss.set_sizes({5, 6, 7});
76+
77+
EXPECT_EQ(ss.size_at_unchecked(0), 5);
78+
EXPECT_EQ(ss.size_at_unchecked(1), 6);
79+
EXPECT_EQ(ss.size_at_unchecked(2), 7);
80+
}
81+
82+
TEST(SizesAndStridesTest, ModifySizeAt) {
83+
SizesAndStrides ss;
84+
ss.set_sizes({2, 3, 4});
85+
86+
ss.size_at(1) = 10;
87+
88+
EXPECT_EQ(ss.size_at(1), 10);
89+
}
90+
91+
TEST(SizesAndStridesTest, ModifyStrideAt) {
92+
SizesAndStrides ss;
93+
ss.set_sizes({2, 3});
94+
ss.set_strides({3, 1});
95+
96+
ss.stride_at(0) = 100;
97+
98+
EXPECT_EQ(ss.stride_at(0), 100);
99+
}
100+
101+
// =============================================================================
102+
// Inline vs Out-of-Line Storage Tests
103+
// =============================================================================
104+
105+
TEST(SizesAndStridesTest, InlineStorage) {
106+
SizesAndStrides ss;
107+
ss.set_sizes({1, 2, 3, 4, 5});
108+
ss.set_strides({120, 60, 20, 5, 1});
109+
110+
EXPECT_EQ(ss.size(), 5u);
111+
EXPECT_EQ(ss.sizes_arrayref()[4], 5);
112+
EXPECT_EQ(ss.strides_arrayref()[4], 1);
113+
}
114+
115+
TEST(SizesAndStridesTest, OutOfLineStorage) {
116+
SizesAndStrides ss;
117+
ss.set_sizes({1, 2, 3, 4, 5, 6, 7});
118+
ss.set_strides({5040, 2520, 840, 210, 42, 7, 1});
119+
120+
EXPECT_EQ(ss.size(), 7u);
121+
EXPECT_EQ(ss.sizes_arrayref()[6], 7);
122+
EXPECT_EQ(ss.strides_arrayref()[6], 1);
123+
}
124+
125+
TEST(SizesAndStridesTest, ResizeFromInlineToOutOfLine) {
126+
SizesAndStrides ss;
127+
ss.set_sizes({1, 2, 3});
128+
129+
EXPECT_EQ(ss.size(), 3u);
130+
131+
ss.set_sizes({1, 2, 3, 4, 5, 6, 7, 8});
132+
ss.set_strides({40320, 20160, 6720, 1680, 336, 56, 8, 1});
133+
134+
EXPECT_EQ(ss.size(), 8u);
135+
EXPECT_EQ(ss.sizes_arrayref()[7], 8);
136+
}
137+
138+
TEST(SizesAndStridesTest, ResizeFromOutOfLineToInline) {
139+
SizesAndStrides ss;
140+
ss.set_sizes({1, 2, 3, 4, 5, 6, 7});
141+
ss.set_strides({5040, 2520, 840, 210, 42, 7, 1});
142+
143+
EXPECT_EQ(ss.size(), 7u);
144+
145+
ss.set_sizes({2, 3});
146+
ss.set_strides({3, 1});
147+
148+
EXPECT_EQ(ss.size(), 2u);
149+
EXPECT_EQ(ss.sizes_arrayref()[0], 2);
150+
EXPECT_EQ(ss.strides_arrayref()[0], 3);
151+
}
152+
153+
// =============================================================================
154+
// Copy and Move Tests
155+
// =============================================================================
156+
157+
TEST(SizesAndStridesTest, CopyConstructInline) {
158+
SizesAndStrides original;
159+
original.set_sizes({2, 3, 4});
160+
original.set_strides({12, 4, 1});
161+
162+
SizesAndStrides copy = original;
163+
164+
EXPECT_EQ(copy.size(), 3u);
165+
EXPECT_EQ(copy.sizes_arrayref()[0], 2);
166+
EXPECT_EQ(copy.strides_arrayref()[0], 12);
167+
}
168+
169+
TEST(SizesAndStridesTest, CopyConstructOutOfLine) {
170+
SizesAndStrides original;
171+
original.set_sizes({1, 2, 3, 4, 5, 6, 7});
172+
original.set_strides({5040, 2520, 840, 210, 42, 7, 1});
173+
174+
SizesAndStrides copy = original;
175+
176+
EXPECT_EQ(copy.size(), 7u);
177+
EXPECT_EQ(copy.sizes_arrayref()[6], 7);
178+
EXPECT_EQ(copy.strides_arrayref()[6], 1);
179+
}
180+
181+
TEST(SizesAndStridesTest, CopyAssignInline) {
182+
SizesAndStrides original;
183+
original.set_sizes({2, 3});
184+
original.set_strides({3, 1});
185+
186+
SizesAndStrides copy;
187+
copy = original;
188+
189+
EXPECT_EQ(copy.size(), 2u);
190+
EXPECT_EQ(copy.sizes_arrayref()[1], 3);
191+
}
192+
193+
TEST(SizesAndStridesTest, MoveConstructInline) {
194+
SizesAndStrides original;
195+
original.set_sizes({2, 3, 4});
196+
original.set_strides({12, 4, 1});
197+
198+
SizesAndStrides moved = std::move(original);
199+
200+
EXPECT_EQ(moved.size(), 3u);
201+
EXPECT_EQ(moved.sizes_arrayref()[0], 2);
202+
EXPECT_EQ(original.size(), 0u);
203+
}
204+
205+
TEST(SizesAndStridesTest, MoveConstructOutOfLine) {
206+
SizesAndStrides original;
207+
original.set_sizes({1, 2, 3, 4, 5, 6, 7});
208+
original.set_strides({5040, 2520, 840, 210, 42, 7, 1});
209+
210+
SizesAndStrides moved = std::move(original);
211+
212+
EXPECT_EQ(moved.size(), 7u);
213+
EXPECT_EQ(moved.sizes_arrayref()[6], 7);
214+
EXPECT_EQ(original.size(), 0u);
215+
}
216+
217+
TEST(SizesAndStridesTest, MoveAssignOutOfLine) {
218+
SizesAndStrides original;
219+
original.set_sizes({1, 2, 3, 4, 5, 6, 7});
220+
original.set_strides({5040, 2520, 840, 210, 42, 7, 1});
221+
222+
SizesAndStrides target;
223+
target = std::move(original);
224+
225+
EXPECT_EQ(target.size(), 7u);
226+
EXPECT_EQ(target.sizes_arrayref()[6], 7);
227+
EXPECT_EQ(original.size(), 0u);
228+
}
229+
230+
// =============================================================================
231+
// Equality Tests
232+
// =============================================================================
233+
234+
TEST(SizesAndStridesTest, EqualityTrue) {
235+
SizesAndStrides ss1;
236+
ss1.set_sizes({2, 3, 4});
237+
ss1.set_strides({12, 4, 1});
238+
239+
SizesAndStrides ss2;
240+
ss2.set_sizes({2, 3, 4});
241+
ss2.set_strides({12, 4, 1});
242+
243+
EXPECT_TRUE(ss1 == ss2);
244+
}
245+
246+
TEST(SizesAndStridesTest, EqualityFalseDifferentSizes) {
247+
SizesAndStrides ss1;
248+
ss1.set_sizes({2, 3, 4});
249+
ss1.set_strides({12, 4, 1});
250+
251+
SizesAndStrides ss2;
252+
ss2.set_sizes({2, 3, 5});
253+
ss2.set_strides({15, 5, 1});
254+
255+
EXPECT_FALSE(ss1 == ss2);
256+
}
257+
258+
TEST(SizesAndStridesTest, EqualityFalseDifferentDims) {
259+
SizesAndStrides ss1;
260+
ss1.set_sizes({2, 3});
261+
262+
SizesAndStrides ss2;
263+
ss2.set_sizes({2, 3, 4});
264+
265+
EXPECT_FALSE(ss1 == ss2);
266+
}
267+
268+
// =============================================================================
269+
// Iterator Tests
270+
// =============================================================================
271+
272+
TEST(SizesAndStridesTest, SizesIterator) {
273+
SizesAndStrides ss;
274+
ss.set_sizes({2, 3, 4});
275+
276+
int64_t sum = 0;
277+
for (auto it = ss.sizes_begin(); it != ss.sizes_end(); ++it) {
278+
sum += *it;
279+
}
280+
EXPECT_EQ(sum, 9);
281+
}
282+
283+
TEST(SizesAndStridesTest, StridesIterator) {
284+
SizesAndStrides ss;
285+
ss.set_sizes({2, 3, 4});
286+
ss.set_strides({12, 4, 1});
287+
288+
int64_t sum = 0;
289+
for (auto it = ss.strides_begin(); it != ss.strides_end(); ++it) {
290+
sum += *it;
291+
}
292+
EXPECT_EQ(sum, 17);
293+
}
294+
295+
// =============================================================================
296+
// Edge Cases
297+
// =============================================================================
298+
299+
TEST(SizesAndStridesTest, ScalarTensor) {
300+
SizesAndStrides ss;
301+
ss.resize(0);
302+
303+
EXPECT_EQ(ss.size(), 0u);
304+
}
305+
306+
TEST(SizesAndStridesTest, OneDimensional) {
307+
SizesAndStrides ss;
308+
ss.set_sizes({10});
309+
ss.set_strides({1});
310+
311+
EXPECT_EQ(ss.size(), 1u);
312+
EXPECT_EQ(ss.size_at(0), 10);
313+
EXPECT_EQ(ss.stride_at(0), 1);
314+
}
315+
316+
} // namespace executorch::backends::aoti::slim::c10

‎backends/aoti/slim/c10/macros/Macros.h‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
// SlimTensor Macros Header
12+
//
13+
// This header provides common compiler hint macros for SlimTensor code.
14+
15+
// Compiler hint macros for branch prediction
16+
#if defined(__GNUC__) || defined(__clang__)
17+
#define SLIMTENSOR_LIKELY(x) __builtin_expect(!!(x), 1)
18+
#define SLIMTENSOR_UNLIKELY(x) __builtin_expect(!!(x), 0)
19+
#else
20+
#define SLIMTENSOR_LIKELY(x) (x)
21+
#define SLIMTENSOR_UNLIKELY(x) (x)
22+
#endif
23+
24+
namespace executorch::backends::aoti::slim::c10 {
25+
// Empty namespace - macros are defined above
26+
} // namespace executorch::backends::aoti::slim::c10

‎backends/aoti/slim/c10/macros/TARGETS‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
load("targets.bzl", "define_common_targets")
2+
3+
define_common_targets()

0 commit comments

Comments
 (0)