1 /*
2 * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25 package hat.buffer;
26
27 import hat.annotations.Kernel;
28 import hat.callgraph.KernelCallGraph;
29 import optkl.ifacemapper.AccessType;
30 import optkl.ifacemapper.BoundSchema;
31 import optkl.ifacemapper.Buffer;
32 import optkl.ifacemapper.Schema;
33 import optkl.ifacemapper.SchemaBuilder;
34 import optkl.util.carriers.ArenaAndLookupCarrier;
35
36 import java.lang.annotation.Annotation;
37 import java.lang.foreign.MemorySegment;
38
39 import static java.lang.foreign.ValueLayout.JAVA_BYTE;
40 import static java.lang.foreign.ValueLayout.JAVA_INT;
41 import static optkl.ifacemapper.MappableIface.getMemorySegment;
42
43 public interface ArgArray extends Buffer {
44 interface Arg extends Buffer.Struct {
45 interface Value extends Buffer.Union {
46 interface Buf extends Buffer.Struct {
47 MemorySegment address();
48
49 void address(MemorySegment address);
50
51 long bytes();
52
53 void bytes(long bytes);
54
55 byte access();
56
57 void access(byte access);
58 }
59
60 boolean z1();
61
62 void z1(boolean z1);
63
64 byte s8();
65
66 void s8(byte s8);
67
68 char u16();
69
70 void u16(char u16);
71
72 short s16();
73
74 void s16(short s16);
75
76 int u32();
77
78 void u32(int u32);
79
80 int s32();
81
82 void s32(int s32);
83
84 float f32();
85
86 void f32(float f32);
87
88 long u64();
89
90 void u64(long u64);
91
92 long s64();
93
94 void s64(long s64);
95
96 double f64();
97
98 void f64(double f64);
99
100 Buf buf();
101 }
102
103 int idx();
104
105 void idx(int idx);
106
107 byte variant();
108
109 void variant(byte variant);
110
111 Value value();
112
113 default String asString() {
114 return switch (variant()) {
115 case '&' -> Long.toHexString(u64());
116 case 'F' -> Float.toString(f32());
117 case 'I' -> Integer.toString(s32());
118 case 'J' -> Long.toString(s64());
119 case 'D' -> Double.toString(f64());
120 case 'Z' -> Boolean.toString(z1());
121 case 'B' -> Byte.toString(s8());
122 case 'S' -> Short.toString(s16());
123 case 'C' -> Character.toString(u16());
124 default -> throw new IllegalStateException("Variant " + variant() + " not supported for arg");
125 };
126 }
127
128 default boolean z1() {
129 return value().z1();
130 }
131
132 default void z1(boolean z1) {
133 variant((byte) 'Z');
134 value().z1(z1);
135 }
136
137 default byte s8() {
138 return value().s8();
139 }
140
141 default void s8(byte s8) {
142 variant((byte) 'B');
143 value().s8(s8);
144 }
145
146 default char u16() {
147 return value().u16();
148 }
149
150 default void u16(char u16) {
151 variant((byte) 'C');
152 value().u16(u16);
153 }
154
155 default short s16() {
156 return value().s16();
157 }
158
159 default void s16(short s16) {
160 variant((byte) 'S');
161 value().s16(s16);
162 }
163
164 default int s32() {
165 return value().s32();
166 }
167
168 default void s32(int s32) {
169 variant((byte) 'I');
170 value().s32(s32);
171 }
172
173 default float f32() {
174 return value().f32();
175 }
176
177 default void f32(float f32) {
178 variant((byte) 'F');
179 value().f32(f32);
180 }
181
182 default long s64() {
183 return value().s64();
184 }
185
186 default void s64(long s64) {
187 variant((byte) 'J');
188 value().s64(s64);
189 }
190
191 default long u64() {
192 return value().u64();
193 }
194
195 default void u64(long u64) {
196 variant((byte) '&');
197 value().u64(u64);
198 }
199
200 default double f64() {
201 return value().f64();
202 }
203
204 default void f64(double f64) {
205 variant((byte) 'D');
206 value().f64(f64);
207 }
208 }
209
210 int argc();
211
212 Arg arg(long idx);
213
214 int schemaLen();
215
216 byte schemaBytes(long idx);
217
218 void schemaBytes(long idx, byte b);
219
220 Schema<ArgArray> schema = Schema.of(ArgArray.class, s -> s
221 .arrayLen("argc")
222 .pad((int) (16 - JAVA_INT.byteSize()))
223 .array("arg", arg -> arg
224 .fields("idx", "variant")
225 .pad((int) (16 - JAVA_INT.byteSize() - JAVA_BYTE.byteSize()))
226 .field("value", val -> val
227 .fields("z1", "s8", "u16", "s16", "s32", "u32", "f32", "s64", "u64", "f64")
228 .field("buf", buf -> buf
229 .fields("address", "bytes", "access")
230 .pad((int) (16 - JAVA_BYTE.byteSize()))
231 )
232 )
233 )
234 .arrayLen("schemaLen")
235 .array("schemaBytes")
236 );
237
238 static ArgArray create(ArenaAndLookupCarrier cc, KernelCallGraph kernelCallGraph, Object[] dispatchContextAndArgs) {
239 String[] schemas = new String[dispatchContextAndArgs.length];
240 StringBuilder argSchema = new StringBuilder();
241 argSchema.append(dispatchContextAndArgs.length);
242 for (int i = 0; i < dispatchContextAndArgs.length; i++) {
243 Object argObject = dispatchContextAndArgs[i];
244 schemas[i] = switch (argObject) {
245 case Boolean _ -> "(?:z1)";
246 case Byte _ -> "(?:s8)";
247 case Short _ -> "(?:s16)";
248 case Character _ -> "(?:u16)";
249 case Float _ -> "(?:f32)";
250 case Integer _ -> "(?:s32)";
251 case Long _ -> "(?:s64)";
252 case Double _ -> "(?:f64)";
253 case Buffer buffer -> "(?:" + SchemaBuilder.schema(buffer) + ")";
254 default ->
255 throw new IllegalStateException("Unexpected value: " + argObject + " Did you pass an interface which is neither a Complete or Incomplete buffer");
256 };
257 if (i > 0) {
258 argSchema.append(",");
259 }
260 argSchema.append(schemas[i]);
261 }
262 String schemaStr = argSchema.toString();
263 ArgArray argArray = BoundSchema.of(cc, schema, dispatchContextAndArgs.length, schemaStr.length() + 1).allocate();
264 byte[] schemaStrBytes = schemaStr.getBytes();
265 for (int i = 0; i < schemaStrBytes.length; i++) {
266 argArray.schemaBytes(i, schemaStrBytes[i]);
267 }
268 argArray.schemaBytes(schemaStrBytes.length, (byte) 0);
269 update(argArray, kernelCallGraph, dispatchContextAndArgs);
270 return argArray;
271 }
272
273 private static void checkAnnotationArg(KernelCallGraph kernelCallGraph, AccessType accessType, Annotation annotation, int argIndex) {
274 if (accessType != AccessType.NA && !isKernelAnnotationPresent(kernelCallGraph)) {
275 // print warning to remove the I/O annotation from the arg parameter
276 IO.println("[WARNING]: Annotation " + annotation + " can be removed from the input parameter " + argIndex);
277 }
278 }
279
280 private static boolean isKernelAnnotationPresent(KernelCallGraph kernelCallGraph) {
281 return kernelCallGraph.callDag.entryPoint.method().getAnnotation(Kernel.class) != null;
282 }
283
284 static void update(ArgArray argArray, KernelCallGraph kernelCallGraph, Object[] dispatchContextAndArgs) {
285 Annotation[][] parameterAnnotations = kernelCallGraph.callDag.entryPoint.method().getParameterAnnotations();
286 var bufferAccessList = kernelCallGraph.bufferAccessList;
287 for (int argIndex = 0; argIndex < dispatchContextAndArgs.length; argIndex++) {
288 final int finalArgIndex = argIndex; // Sigh
289 Object argObject = dispatchContextAndArgs[argIndex];
290 Arg arg = argArray.arg(argIndex); // this should be invariant, but if we are called from create it will be 0 for all
291 arg.idx(argIndex);
292 switch (argObject) {
293 case Boolean z1 -> arg.z1(z1);
294 case Byte s8 -> arg.s8(s8);
295 case Short s16 -> arg.s16(s16);
296 case Character u16 -> arg.u16(u16);
297 case Float f32 -> arg.f32(f32);
298 case Integer s32 -> arg.s32(s32);
299 case Long s64 -> arg.s64(s64);
300 case Double f64 -> arg.f64(f64);
301 case Buffer buffer when finalArgIndex >0-> {
302
303 Annotation[] annotations = parameterAnnotations[finalArgIndex-1];
304 AccessType accessType = AccessType.NA;
305 for (Annotation annotation : annotations) {
306 accessType = AccessType.of(annotation);
307 checkAnnotationArg(kernelCallGraph, accessType, annotation, finalArgIndex-1);
308 }
309
310 MemorySegment segment = getMemorySegment(buffer);
311 arg.variant((byte) '&');
312 Arg.Value value = arg.value();
313 Arg.Value.Buf buf = value.buf();
314 buf.address(segment);
315 buf.bytes(segment.byteSize());
316
317 if (isKernelAnnotationPresent(kernelCallGraph)) {
318 // If the annotation is present, then we keep the accessor defined for each parameter
319 buf.access(accessType.value);
320 } else {
321 // otherwise, we rely on the buffer-tagger to set the accessor
322 buf.access(bufferAccessList.get(finalArgIndex-1).value);
323 }
324 }
325 case Buffer buffer -> {
326 MemorySegment segment = getMemorySegment(buffer);
327 arg.variant((byte) '&');
328 Arg.Value value = arg.value();
329 Arg.Value.Buf buf = value.buf();
330 buf.address(segment);
331 buf.bytes(segment.byteSize());
332 }
333 default -> throw new IllegalStateException("Unexpected value: " + argObject+ ":"+argObject.getClass().getName()+" in dispatchAndArgs slot "+argIndex);
334 }
335 }
336 }
337
338 default String getSchemaBytes() {
339 byte[] bytes = new byte[schemaLen() + 1];
340 for (int i = 0; i < schemaLen(); i++) {
341 bytes[i] = schemaBytes(i);
342 }
343 bytes[bytes.length - 1] = '0';
344 return new String(bytes);
345 }
346
347 default String dump() {
348 StringBuilder dump = new StringBuilder();
349 dump.append("SchemaBytes:").append(getSchemaBytes()).append("\n");
350 for (int argIndex = 0; argIndex < argc(); argIndex++) {
351 Arg arg = arg(argIndex);
352 dump.append(arg.asString()).append("\n");
353 }
354 return dump.toString();
355 }
356 }