xref: /aosp_15_r20/external/pytorch/c10/util/AlignOf.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 //===--- AlignOf.h - Portable calculation of type alignment -----*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file defines the AlignedCharArray and AlignedCharArrayUnion classes.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 // ATen: modified from llvm::AlignOf
15 // replaced LLVM_ALIGNAS with alignas
16 
17 #pragma once
18 
19 #include <cstddef>
20 
21 namespace c10 {
22 
23 /// \struct AlignedCharArray
24 /// \brief Helper for building an aligned character array type.
25 ///
26 /// This template is used to explicitly build up a collection of aligned
27 /// character array types. We have to build these up using a macro and explicit
28 /// specialization to cope with MSVC (at least till 2015) where only an
29 /// integer literal can be used to specify an alignment constraint. Once built
30 /// up here, we can then begin to indirect between these using normal C++
31 /// template parameters.
32 
33 // MSVC requires special handling here.
34 #ifndef _MSC_VER
35 
36 template <size_t Alignment, size_t Size>
37 struct AlignedCharArray {
38   // NOLINTNEXTLINE(*c-arrays)
39   alignas(Alignment) char buffer[Size];
40 };
41 
42 #else // _MSC_VER
43 
44 /// \brief Create a type with an aligned char buffer.
45 template <size_t Alignment, size_t Size>
46 struct AlignedCharArray;
47 
48 // We provide special variations of this template for the most common
49 // alignments because __declspec(align(...)) doesn't actually work when it is
50 // a member of a by-value function argument in MSVC, even if the alignment
51 // request is something reasonably like 8-byte or 16-byte. Note that we can't
52 // even include the declspec with the union that forces the alignment because
53 // MSVC warns on the existence of the declspec despite the union member forcing
54 // proper alignment.
55 
56 template <size_t Size>
57 struct AlignedCharArray<1, Size> {
58   union {
59     char aligned;
60     char buffer[Size];
61   };
62 };
63 
64 template <size_t Size>
65 struct AlignedCharArray<2, Size> {
66   union {
67     short aligned;
68     char buffer[Size];
69   };
70 };
71 
72 template <size_t Size>
73 struct AlignedCharArray<4, Size> {
74   union {
75     int aligned;
76     char buffer[Size];
77   };
78 };
79 
80 template <size_t Size>
81 struct AlignedCharArray<8, Size> {
82   union {
83     double aligned;
84     char buffer[Size];
85   };
86 };
87 
88 // The rest of these are provided with a __declspec(align(...)) and we simply
89 // can't pass them by-value as function arguments on MSVC.
90 
91 #define AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT(x) \
92   template <size_t Size>                          \
93   struct AlignedCharArray<x, Size> {              \
94     __declspec(align(x)) char buffer[Size];       \
95   };
96 
97 AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT(16)
98 AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT(32)
99 AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT(64)
100 AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT(128)
101 
102 #undef AT_ALIGNEDCHARARRAY_TEMPLATE_ALIGNMENT
103 
104 #endif // _MSC_VER
105 
106 namespace detail {
107 template <
108     typename T1,
109     typename T2 = char,
110     typename T3 = char,
111     typename T4 = char,
112     typename T5 = char,
113     typename T6 = char,
114     typename T7 = char,
115     typename T8 = char,
116     typename T9 = char,
117     typename T10 = char>
118 class AlignerImpl {
119   T1 t1;
120   T2 t2;
121   T3 t3;
122   T4 t4;
123   T5 t5;
124   T6 t6;
125   T7 t7;
126   T8 t8;
127   T9 t9;
128   T10 t10;
129 
130  public:
131   AlignerImpl() = delete;
132 };
133 
134 template <
135     typename T1,
136     typename T2 = char,
137     typename T3 = char,
138     typename T4 = char,
139     typename T5 = char,
140     typename T6 = char,
141     typename T7 = char,
142     typename T8 = char,
143     typename T9 = char,
144     typename T10 = char>
145 union SizerImpl {
146   // NOLINTNEXTLINE(*c-arrays)
147   char arr1[sizeof(T1)], arr2[sizeof(T2)], arr3[sizeof(T3)], arr4[sizeof(T4)],
148       arr5[sizeof(T5)], arr6[sizeof(T6)], arr7[sizeof(T7)], arr8[sizeof(T8)],
149       arr9[sizeof(T9)], arr10[sizeof(T10)];
150 };
151 } // end namespace detail
152 
153 /// \brief This union template exposes a suitably aligned and sized character
154 /// array member which can hold elements of any of up to ten types.
155 ///
156 /// These types may be arrays, structs, or any other types. The goal is to
157 /// expose a char array buffer member which can be used as suitable storage for
158 /// a placement new of any of these types. Support for more than ten types can
159 /// be added at the cost of more boilerplate.
160 template <
161     typename T1,
162     typename T2 = char,
163     typename T3 = char,
164     typename T4 = char,
165     typename T5 = char,
166     typename T6 = char,
167     typename T7 = char,
168     typename T8 = char,
169     typename T9 = char,
170     typename T10 = char>
171 struct AlignedCharArrayUnion
172     : AlignedCharArray<
173           alignof(detail::AlignerImpl<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>),
174           sizeof(::c10::detail::
175                      SizerImpl<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>)> {};
176 } // end namespace c10
177