1 /*
2 * Copyright (c) 2021 Arm Limited.
3 *
4 * SPDX-License-Identifier: MIT
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in
14 * all copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
21 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
22 * IN THE SOFTWARE.
23 */
24
25 #pragma once
26
27 #ifdef __aarch64__
28
29 namespace {
30
a64_transpose_interleave_16_2x2(uint16_t * out,const uint16_t * in,size_t width,size_t in_stride,size_t height)31 void a64_transpose_interleave_16_2x2(uint16_t *out, const uint16_t *in, size_t width, size_t in_stride, size_t height)
32 {
33 uint16_t *pad_row = reinterpret_cast<uint16_t *>(alloca(width * sizeof(uint16_t)));
34
35 if (height % 2) {
36 memset(pad_row, 0, width * sizeof(uint16_t));
37 }
38
39 size_t out_stride = 16 * roundup<size_t>(height, 2) * sizeof(uint16_t);
40
41 __asm__ __volatile__(
42 "cmp %x[height], #0x8\n"
43 "blt 8f\n"
44 "1:" // Main row loop: Head
45 "mov x28, %x[in]\n"
46 "mov x27, %x[out]\n"
47 "add x26, x28, %x[in_stride]\n"
48 "add x25, x26, %x[in_stride]\n"
49 "add x24, x25, %x[in_stride]\n"
50 "add x23, x24, %x[in_stride]\n"
51 "add x22, x23, %x[in_stride]\n"
52 "add x21, x22, %x[in_stride]\n"
53 "add x20, x21, %x[in_stride]\n"
54 "add %x[in], x20, %x[in_stride]\n"
55 "sub %x[height], %x[height], #0x8\n"
56 "mov x19, %x[width]\n"
57 "cmp x19, #0x10\n"
58 "blt 3f\n"
59 "2:" // Main row loop: Column loop
60 "ldr q18, [x28], #0x10\n"
61 "sub x19, x19, #0x10\n"
62 "ldr q16, [x26], #0x10\n"
63 "zip1 v1.8h, v18.8h, v16.8h\n"
64 "ldr q17, [x28], #0x10\n"
65 "cmp x19, #0x10\n"
66 "zip2 v0.8h, v18.8h, v16.8h\n"
67 "ldr q16, [x26], #0x10\n"
68 "ldr q19, [x25], #0x10\n"
69 "zip1 v31.8h, v17.8h, v16.8h\n"
70 "ldr q18, [x25], #0x10\n"
71 "zip2 v30.8h, v17.8h, v16.8h\n"
72 "ldr q16, [x24], #0x10\n"
73 "ldr q20, [x23], #0x10\n"
74 "zip1 v29.8h, v19.8h, v16.8h\n"
75 "ldr q17, [x24], #0x10\n"
76 "zip2 v28.8h, v19.8h, v16.8h\n"
77 "ldr q19, [x23], #0x10\n"
78 "ldr q16, [x22], #0x10\n"
79 "zip1 v27.8h, v18.8h, v17.8h\n"
80 "ldr q26, [x21], #0x10\n"
81 "zip2 v25.8h, v18.8h, v17.8h\n"
82 "ldr q18, [x22], #0x10\n"
83 "zip1 v24.8h, v20.8h, v16.8h\n"
84 "ldr q17, [x20], #0x10\n"
85 "zip2 v23.8h, v20.8h, v16.8h\n"
86 "ldr q22, [x21], #0x10\n"
87 "zip1 v21.8h, v19.8h, v18.8h\n"
88 "ldr q16, [x20], #0x10\n"
89 "zip2 v20.8h, v19.8h, v18.8h\n"
90 "zip1 v19.8h, v26.8h, v17.8h\n"
91 "str q1, [x27, #0x0]\n"
92 "zip2 v18.8h, v26.8h, v17.8h\n"
93 "str q0, [x27, #0x10]\n"
94 "str q31, [x27, #0x20]\n"
95 "zip1 v17.8h, v22.8h, v16.8h\n"
96 "str q30, [x27, #0x30]\n"
97 "zip2 v16.8h, v22.8h, v16.8h\n"
98 "str q29, [x27, #0x40]\n"
99 "str q28, [x27, #0x50]\n"
100 "str q27, [x27, #0x60]\n"
101 "str q25, [x27, #0x70]\n"
102 "str q24, [x27, #0x80]\n"
103 "str q23, [x27, #0x90]\n"
104 "str q21, [x27, #0xa0]\n"
105 "str q20, [x27, #0xb0]\n"
106 "str q19, [x27, #0xc0]\n"
107 "str q18, [x27, #0xd0]\n"
108 "str q17, [x27, #0xe0]\n"
109 "str q16, [x27, #0xf0]\n"
110 "add x27, x27, %x[out_stride]\n"
111 "bge 2b\n"
112 "3:" // Main row loop: Column loop skip
113 "cmp x19, #0x4\n"
114 "blt 5f\n"
115 "4:" // Main row loop: width 4 loop: loop
116 "ldr d17, [x28], #0x8\n"
117 "sub x19, x19, #0x4\n"
118 "ldr d16, [x26], #0x8\n"
119 "zip1 v20.8h, v17.8h, v16.8h\n"
120 "ldr d17, [x25], #0x8\n"
121 "cmp x19, #0x4\n"
122 "ldr d16, [x24], #0x8\n"
123 "zip1 v19.8h, v17.8h, v16.8h\n"
124 "ldr d17, [x23], #0x8\n"
125 "ldr d16, [x22], #0x8\n"
126 "zip1 v18.8h, v17.8h, v16.8h\n"
127 "ldr d17, [x21], #0x8\n"
128 "ldr d16, [x20], #0x8\n"
129 "zip1 v16.8h, v17.8h, v16.8h\n"
130 "str q20, [x27, #0x0]\n"
131 "str q19, [x27, #0x40]\n"
132 "str q18, [x27, #0x80]\n"
133 "str q16, [x27, #0xc0]\n"
134 "add x27, x27, #0x10\n"
135 "bge 4b\n"
136 "5:" // Main row loop: width 4 loop: skip
137 "cmp x19, #0x1\n"
138 "blt 7f\n"
139 "6:" // Main row loop: width 1 loop: loop
140 "ldr h17, [x28], #0x2\n"
141 "sub x19, x19, #0x1\n"
142 "ldr h16, [x26], #0x2\n"
143 "zip1 v20.8h, v17.8h, v16.8h\n"
144 "ldr h17, [x25], #0x2\n"
145 "cmp x19, #0x1\n"
146 "ldr h16, [x24], #0x2\n"
147 "zip1 v19.8h, v17.8h, v16.8h\n"
148 "ldr h17, [x23], #0x2\n"
149 "ldr h16, [x22], #0x2\n"
150 "zip1 v18.8h, v17.8h, v16.8h\n"
151 "ldr h17, [x21], #0x2\n"
152 "ldr h16, [x20], #0x2\n"
153 "zip1 v16.8h, v17.8h, v16.8h\n"
154 "str s20, [x27, #0x0]\n"
155 "str s19, [x27, #0x40]\n"
156 "str s18, [x27, #0x80]\n"
157 "str s16, [x27, #0xc0]\n"
158 "add x27, x27, #0x4\n"
159 "bge 6b\n"
160 "7:" // Main row loop: width 1 loop: skip
161 "add %x[out], %x[out], #0x100\n"
162 "cmp %x[height], #0x8\n"
163 "bge 1b\n"
164 "cbz %x[height], 16f\n"
165 "8:" // Main loop skip
166
167 "9:" // Tail row loop: Head
168 "mov x28, %x[in]\n"
169 "mov x27, %x[out]\n"
170 "add x26, x28, %x[in_stride]\n"
171 "add %x[in], x26, %x[in_stride]\n"
172 "cmp %x[height], #0x1\n"
173 "csel x26, x26, %x[pad_row], GT\n"
174 "sub %x[height], %x[height], #0x2\n"
175 "mov x19, %x[width]\n"
176 "cmp x19, #0x10\n"
177 "blt 11f\n"
178 "10:" // Tail row loop: Column loop
179 "ldr q18, [x28], #0x10\n"
180 "sub x19, x19, #0x10\n"
181 "ldr q16, [x26], #0x10\n"
182 "zip1 v17.8h, v18.8h, v16.8h\n"
183 "ldr q19, [x28], #0x10\n"
184 "cmp x19, #0x10\n"
185 "zip2 v18.8h, v18.8h, v16.8h\n"
186 "ldr q16, [x26], #0x10\n"
187 "str q17, [x27, #0x0]\n"
188 "zip1 v17.8h, v19.8h, v16.8h\n"
189 "str q18, [x27, #0x10]\n"
190 "zip2 v16.8h, v19.8h, v16.8h\n"
191 "str q17, [x27, #0x20]\n"
192 "str q16, [x27, #0x30]\n"
193 "add x27, x27, %x[out_stride]\n"
194 "bge 10b\n"
195 "11:" // Tail row loop: Column loop skip
196 "cmp x19, #0x4\n"
197 "blt 13f\n"
198 "12:" // Tail row loop: width 4 loop: loop
199 "ldr d17, [x28], #0x8\n"
200 "sub x19, x19, #0x4\n"
201 "ldr d16, [x26], #0x8\n"
202 "zip1 v16.8h, v17.8h, v16.8h\n"
203 "str q16, [x27, #0x0]\n"
204 "add x27, x27, #0x10\n"
205 "cmp x19, #0x4\n"
206 "bge 12b\n"
207 "13:" // Tail row loop: width 4 loop: skip
208 "cmp x19, #0x1\n"
209 "blt 15f\n"
210 "14:" // Tail row loop: width 1 loop: loop
211 "ldr h17, [x28], #0x2\n"
212 "sub x19, x19, #0x1\n"
213 "ldr h16, [x26], #0x2\n"
214 "zip1 v16.8h, v17.8h, v16.8h\n"
215 "str s16, [x27, #0x0]\n"
216 "add x27, x27, #0x4\n"
217 "cmp x19, #0x1\n"
218 "bge 14b\n"
219 "15:" // Tail row loop: width 1 loop: skip
220 "add %x[out], %x[out], #0x40\n"
221 "cmp %x[height], #0x1\n"
222 "bge 9b\n"
223 "16:" // Done
224
225 : [height] "+&r" (height), [in] "+&r" (in), [out] "+&r" (out)
226 : [in_stride] "r" (in_stride), [out_stride] "r" (out_stride), [pad_row] "r" (pad_row), [width] "r" (width)
227 : "cc", "memory", "v0", "v1", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "x19", "x20", "x21", "x22", "x23", "x24", "x25", "x26", "x27", "x28"
228 );
229 }
230
231 } // anonymous namespace
232
233 template<>
Transform(bfloat16 * out,const bfloat16 * in,int stride,int x0,int xmax,int k0,int kmax)234 void Transform<16, 2, true, VLType::None>(
235 bfloat16 *out, const bfloat16 *in, int stride, int x0, int xmax, int k0, int kmax)
236 {
237 a64_transpose_interleave_16_2x2(
238 reinterpret_cast<uint16_t *>(out),
239 reinterpret_cast<const uint16_t *>(in + k0 * stride + x0),
240 (xmax-x0) * sizeof(bfloat16) / 2,
241 stride * sizeof(bfloat16),
242 (kmax-k0)
243 );
244 }
245
246 #endif
247