1 /*
2 * Copyright (c) 2024-2026, 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.codebuilders;
26
27 import hat.buffer.Uniforms;
28 import hat.types.F32;
29 import hat.types.ivec2;
30 import hat.types.mat2;
31 import hat.types.mat3;
32 import hat.types.vec2;
33 import hat.types.vec3;
34 import hat.types.vec4;
35 import jdk.incubator.code.Op;
36 import jdk.incubator.code.dialect.java.ClassType;
37 import jdk.incubator.code.dialect.java.JavaType;
38 import jdk.incubator.code.dialect.java.PrimitiveType;
39 import optkl.IfaceValue;
40 import optkl.OpHelper;
41 import optkl.ifacemapper.MappableIface;
42
43 import java.lang.invoke.MethodHandles;
44 import java.lang.reflect.Field;
45 import java.util.Collections;
46 import java.util.List;
47 import java.util.function.Consumer;
48 import java.util.stream.Stream;
49
50 public class C99VecAndMatHandler {
51
52 static final String SHADER_MAIN_IMAGE = "mainImage";
53 static final String SHAPE_FIELD_NAME = "shape";
54
55 static IfaceValue.vec.Shape getVecShape(java.lang.reflect.Type vecClass) {
56 try {
57 Field field = ((Class<?>) vecClass).getField(SHAPE_FIELD_NAME);
58 return (IfaceValue.vec.Shape) field.get(null);
59 } catch (Throwable t) {
60 throw new RuntimeException(t);
61 }
62 }
63
64 static IfaceValue.vec.Shape getVecShape(MethodHandles.Lookup lookup, JavaType javaType) {
65 var resolved = OpHelper.classTypeToTypeOrThrow(lookup, (ClassType) javaType);
66 return getVecShape(resolved);
67
68 }
69
70 static String clName(MethodHandles.Lookup lookup, JavaType javaType) {
71
72 if (OpHelper.isAssignable(lookup, javaType, vec4.class)) {
73 return "vec4";
74 } else if (OpHelper.isAssignable(lookup, javaType, vec3.class)) {
75 return "vec3";
76 } else if (OpHelper.isAssignable(lookup, javaType, vec2.class)) {
77 return "vec2";
78 } else if (OpHelper.isAssignable(lookup, javaType, mat2.class)) {
79 return "mat2";
80 } else if (OpHelper.isAssignable(lookup, javaType, mat3.class)) {
81 return "mat3";
82 } else if (OpHelper.isAssignable(lookup, javaType, ivec2.class)) {
83 return "ivec2";
84 } else if (javaType.equals(PrimitiveType.FLOAT)) {
85 return "float";
86 } else if (javaType.equals(PrimitiveType.INT)) {
87 return "int";
88 } else {
89 throw new RuntimeException("no cl name mapping for " + javaType);
90 }
91 }
92
93 public static boolean isVecInvoke(OpHelper.Invoke invoke) {
94 return (invoke.named(SHADER_MAIN_IMAGE) || invoke.refIs(F32.class, Uniforms.class, IfaceValue.vec.class, IfaceValue.mat.class));
95 }
96
97 public static String mangledName(OpHelper.Invoke invoke) {
98 // So invoke float mod(float lhs, float rhs) -> f32_mod_f32_f32"
99 return clName(invoke.lookup(), (JavaType) invoke.returnType()) + "_"
100 + invoke.name() + "_"
101 + clName(invoke.lookup(), (JavaType) invoke.resultFromOperandNOrNull(0).type()) + "_"
102 + clName(invoke.lookup(), (JavaType) invoke.resultFromOperandNOrNull(1).type());
103 }
104
105 public static <T extends C99HATKernelBuilder<T>> void args(C99HATKernelBuilder<T> bldr, Stream<Op.Result> argStream) {
106 bldr.paren(_ -> bldr.commaSpaceSeparated(argStream, operand -> bldr.recurse(operand.op())));
107 }
108
109 public static <T extends C99HATKernelBuilder<T>> void nameAndArgs(String name, C99HATKernelBuilder<T> bldr, Stream<Op.Result> argStream) {
110 bldr.funcName(name).paren(_ -> args(bldr, argStream));
111 }
112
113 public static <T extends C99HATKernelBuilder<T>> void nameAndArgs(OpHelper.Invoke invoke, String newName, C99HATKernelBuilder<T> bldr) {
114 bldr.funcName(newName).paren(_ ->
115 bldr.commaSpaceSeparated(invoke.operandsAsResults(), operand -> bldr.recurse(operand.op()))
116 );
117 }
118
119 public static <T extends C99HATKernelBuilder<T>> void nameAndArgs(OpHelper.Invoke invoke, C99HATKernelBuilder<T> bldr) {
120 nameAndArgs(invoke, invoke.name(), bldr);
121 }
122
123 public static <T extends C99HATKernelBuilder<T>> void mangledNameAndArgs(OpHelper.Invoke invoke, C99HATKernelBuilder<T> bldr) {
124 nameAndArgs(invoke, mangledName(invoke), bldr);
125 }
126
127 public static boolean hasVecOperand(OpHelper.Invoke invoke) {
128 return invoke.operandsAsResults().anyMatch(r ->
129 r.type() instanceof ClassType classType
130 && OpHelper.classTypeToTypeOrThrow(invoke.lookup(), classType) instanceof Class<?> clazz
131 && IfaceValue.mat.class.isAssignableFrom(clazz));
132
133 }
134
135 public static <T extends C99HATKernelBuilder<T>> void handleF32Invoke(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
136 switch (invoke.name()) {
137 case "fract" -> //opencl 's fract wants a ptr as the second arg where it returns fract value. So we implement our own
138 bldr.paren(_ -> bldr.recurse(invoke.opFromFirstOperandOrNull()).sub().funcName("floor")
139 .paren(_ -> bldr.recurse(invoke.opFromFirstOperandOrNull())));
140 case "atan" -> {
141 if (invoke.operandCount() > 1) { // atan -> atan2
142 nameAndArgs(invoke, invoke.name() + "2", bldr); // atan(l,r) ->atan2(l,r)
143 } else {
144 nameAndArgs(invoke, bldr); // atan(v)->atan(v)
145 }
146 }
147 case "inversesqrt" -> nameAndArgs(invoke, "rsqrt", bldr);// inversesqrt(...)->rsqrt(...)
148 case "cos", "sqrt", "sin", "exp", "pow", "min", "max", "log", "smoothstep", "clamp", "floor", "step",
149 "mix" -> nameAndArgs(invoke, bldr); // asis!
150 case "abs" -> nameAndArgs(invoke, "fabs", bldr);// abs(...)->fabs(...)
151 case "mod" -> mangledNameAndArgs(invoke, bldr); //mod(...) -> float_mod_float_float(...)
152 case "round" -> handleRoundInvoke(bldr, invoke);
153 default -> throw new RuntimeException("unmapped F32 call " + invoke.name());
154 }
155 }
156
157 public static <T extends C99HATKernelBuilder<T>> void handleUniformsInvoke(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
158 bldr.cast(_ -> bldr.type((JavaType) invoke.returnType()));
159 switch (invoke.name()) {
160 case "iResolution" -> bldr.paren(_ ->
161 bldr.sep(List.of("x", "y", "z"), _ -> bldr.csp(), lane -> bldr.id("uniforms").rarrow().id(invoke.name()).dot().id(lane))
162 );
163 case "iMouse" ->
164 bldr.paren(_ -> bldr.sep(List.of("x", "y"), _ -> bldr.csp(), lane -> bldr.id("uniforms").rarrow().id(invoke.name()).dot().id(lane))
165 );
166 case "iTime" -> bldr.id("uniforms").rarrow().id(invoke.name());
167 default -> throw new RuntimeException("some other uniform" + invoke.name());
168 }
169 }
170
171 public static <T extends C99HATKernelBuilder<T>> void handleVecInvoke(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
172 if (invoke.refIs(IfaceValue.Struct.class)) {
173 // A call on a field of the uniforms
174 // so uniforms_t.iTime -> uniforms.iTime;
175 bldr.recurse(invoke.opFromFirstOperandOrNull()).dot().id(invoke.name());
176 } else if (invoke instanceof OpHelper.Invoke.Virtual && invoke.operandCount() == 1 && invoke.refIs(vec3.class, vec2.class, vec4.class)) {
177 // an accessor on a vec say v.x() -> v.x
178 bldr.recurse(invoke.opFromFirstOperandOrNull()).dot().id(invoke.name());
179 } else if (invoke.nameMatchesRegex("vec[234]")) { // a psuedo vec2 constructor vec2(....) -> (vec2)(....)
180 bldr.paren(_ -> bldr.type(invoke.name()));
181 args(bldr, invoke.operandsAsResults());
182 } else if (invoke.nameMatchesRegex("(mod|reflect)")) {
183 mangledNameAndArgs(invoke, bldr);
184 } else if (invoke.named("mul") && invoke.operandCount() == 2 && hasVecOperand(invoke)) {
185 mangledNameAndArgs(invoke, bldr);
186 } else if (invoke.nameMatchesRegex("(mul|add|sub|div)")) {
187 // for opencl we can turn these into expressions. So vec3.mul(l,r) -> (l * r)
188 bldr.paren(_ -> bldr.paren(_ -> bldr.recurse(invoke.opFromFirstOperandOrNull())).symbol(switch (invoke.name()) {
189 case "mul" -> "*";
190 case "add" -> "+";
191 case "div" -> "/";
192 case "sub" -> "-";
193 default -> throw new IllegalStateException("oh my");
194 }).paren(_ -> bldr.recurse(invoke.opFromOperandNOrNull(1))));
195 } else if (invoke.named("abs") && invoke.refIs(IfaceValue.vec.class)) {
196 bldr.funcName("f" + invoke.name()).paren(_ ->
197 bldr.sep(invoke.op().operands(), _ -> bldr.csp(), v -> bldr.recurse(v.asResult().op()))
198 );
199 } else if (invoke.named("fract") && invoke.operandCount() == 1) {
200 // return x - floor(x);
201 bldr.paren(_ -> bldr.recurse(invoke.opFromFirstOperandOrNull()).sub().funcName("floor")
202 .paren(_ -> bldr.recurse(invoke.opFromFirstOperandOrNull())));
203 } else if (invoke.named("round")) {
204 handleRoundInvoke(bldr, invoke);
205 } else if (invoke.nameMatchesRegex("(dot|length|max|mix|min|smoothstep|step|normalize|clamp|pow|cross|distance|floor|fract|sin|cos|abs)")) {
206 bldr.funcName(invoke.op()).paren(_ ->
207 bldr.sep(invoke.op().operands(), _ -> bldr.csp(), v -> bldr.recurse(v.asResult().op()))
208 );
209 } else {
210 StringBuilder stringBuilder = new StringBuilder("For vec types we need to IMPLEMENT " + invoke.refType() + ":" + invoke.name() + "(");
211 invoke.op().operands().forEach(o -> stringBuilder.append(" " + o.asResult().type()));
212 stringBuilder.append(")");
213 throw new RuntimeException(stringBuilder.toString());
214 }
215 }
216
217 public static <T extends C99HATKernelBuilder<T>> void handleMatInvoke(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
218 if (invoke.nameMatchesRegex("mat[23]")) {
219 handleMatConstructor(bldr, invoke);
220 } else if (invoke instanceof OpHelper.Invoke.Virtual && invoke.operandCount() == 1
221 && invoke.nameMatchesRegex("_[0-2][0-2]")) {
222 bldr.recurse(invoke.opFromFirstOperandOrNull()).dot().id(invoke.name());
223 } else if (invoke.named("mul")) {
224 mangledNameAndArgs(invoke, bldr);
225 } else {
226 StringBuilder stringBuilder = new StringBuilder("For mat types we need to IMPLEMENT " + invoke.refType() + ":" + invoke.name() + "(");
227 invoke.op().operands().forEach(o -> stringBuilder.append(" " + o.asResult().type()));
228 stringBuilder.append(")");
229 throw new RuntimeException(stringBuilder.toString());
230 }
231 }
232
233 public static <T extends C99HATKernelBuilder<T>> void handleMatConstructor(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
234 if (invoke.operandCount() == 1 && invoke.resultFromOperandNOrNull(0).type() instanceof ClassType) {
235 bldr.recurse(invoke.opFromFirstOperandOrNull());
236 return;
237 }
238 if (invoke.operandCount() == 1) {
239 nameAndArgs(invoke, invoke.name() + "_splat", bldr);
240 }else {
241 int lanes = invoke.named("mat2") ? 4 : 9;
242 if (invoke.operandCount() != lanes) {
243 throw new UnsupportedOperationException("Unsupported shader matrix constructor " + invoke.name());
244 }
245 // Maybe bldr.type(invoke.name());
246 bldr.paren(_ -> bldr.type(invoke.name()));
247 bldr.brace(_ -> bldr.commaSpaceSeparated(invoke.operandsAsResults(), operand -> bldr.recurse(operand.op())));
248 }
249 }
250
251 public static <T extends C99HATKernelBuilder<T>> void handleInvoke(C99HATKernelBuilder<T> bldr, OpHelper.Invoke invoke) {
252 if (invoke.named(SHADER_MAIN_IMAGE)) {
253 nameAndArgs(invoke, bldr);
254 } else if (invoke.refIs(F32.class)) {
255 handleF32Invoke(bldr, invoke);
256 } else if (invoke.refIs(Uniforms.class)) {
257 handleUniformsInvoke(bldr, invoke);
258 } else if (invoke.refIs(IfaceValue.vec.class)) {
259 handleVecInvoke(bldr, invoke);
260 } else if (invoke.refIs(IfaceValue.mat.class)) {
261 handleMatInvoke(bldr, invoke);
262 }
263 }
264
265 public static boolean isVecOrMatType(MethodHandles.Lookup lookup, JavaType javaType) {
266 return OpHelper.isAssignable(lookup, javaType, MappableIface.vec.class, IfaceValue.mat.class);
267 }
268
269 public static <T extends C99HATKernelBuilder<T>> void handleType(C99HATKernelBuilder<T> bldr, JavaType javaType) {
270 bldr.type(clName(bldr.scopedCodeBuilderContext().lookup(), javaType));
271 }
272
273 public static <T extends C99HATKernelBuilder<T>> void genFunc(C99HATKernelBuilder<T> bldr, String ret, String op, String lhs, String rhs, Consumer<C99HATKernelBuilder<T>> consumer) {
274 bldr.func(
275 _ -> bldr.type(ret),
276 ret + "_" + op + "_" + lhs + "_" + rhs,
277 _ -> bldr.type(lhs).sp().id("l").csp().type(rhs).sp().id("r"),
278 _ -> consumer.accept(bldr)
279 ).semicolon();
280
281 }
282
283 public static <T extends C99HATKernelBuilder<T>> void createVecFunctions(C99HATKernelBuilder<T> builder) {
284 record NamedVecShape(String name, IfaceValue.vec.Shape shape) {
285 }
286
287 List.of(new NamedVecShape("vec2", vec2.shape), new NamedVecShape("vec3", vec3.shape), new NamedVecShape("vec4", vec4.shape)).forEach(ns ->
288 builder.typedefKeyword().sp().type("float" + ns.shape.lanes()).sp().type(ns.name).snl()
289 );
290 List.of(new NamedVecShape("ivec2", vec2.shape)).forEach(ns ->
291 builder.typedefKeyword().sp().type("int" + ns.shape.lanes()).sp().type(ns.name).snl()
292 );
293 createMatFunctions(builder, false);
294
295 createModFunctions(builder);
296 createRoundFunctions(builder, false);
297
298 genFunc(builder, "vec3", "reflect", "vec3", "vec3", _ ->
299 builder.returnKeyword().sp().id("l").sp().minus().id("r").sp().mul().sp()
300 .funcName("dot").paren(_ -> builder.id("r").csp().id("l"))
301 .sp().mul().floatConst(2).semicolon());
302 }
303
304 public static <T extends C99HATKernelBuilder<T>> void createMatFunctions(C99HATKernelBuilder<T> builder, boolean metal) {
305 record NamedMatShape(String name, IfaceValue.mat.Shape shape) {
306 }
307 List.of(new NamedMatShape("mat2", mat2.shape), new NamedMatShape("mat3", mat3.shape)).forEach(ns -> {
308 builder.typedefKeyword().sp().structKeyword().sp().id(ns.name + "_s").braceNlIndented(_ -> {
309 builder.sep(ns.shape.rowColNames(), _ -> builder.nl(), n ->
310 builder.type((JavaType) ns.shape.codeType()).sp().id(n).semicolon());
311 }).sp().id(ns.name).snl();
312 String elements = String.join(",", Collections.nCopies(ns.name.equals("mat2") ? 4 : 9, "s"));
313 builder.func(_ -> builder.type(ns.name), ns.name + "_splat",
314 _ -> builder.type("float").sp().id("s"),
315 _ -> matrixFunctionReturn(builder, ns.name, elements, metal)).semicolon();
316 });
317
318 // Matrix elements are ordered by column, then row; vectors on the left
319 // are row vectors and vectors on the right are column vectors.
320 genFunc(builder, "vec2", "mul", "vec2", "mat2", _ ->
321 matrixFunctionReturn(builder, "vec2", "l.x*r._00+l.y*r._01,l.x*r._10+l.y*r._11", metal));
322 genFunc(builder, "vec2", "mul", "mat2", "vec2", _ ->
323 matrixFunctionReturn(builder, "vec2", "l._00*r.x+l._10*r.y,l._01*r.x+l._11*r.y", metal));
324 genFunc(builder, "vec3", "mul", "vec3", "mat3", _ ->
325 matrixFunctionReturn(builder, "vec3",
326 "l.x*r._00+l.y*r._01+l.z*r._02,l.x*r._10+l.y*r._11+l.z*r._12,l.x*r._20+l.y*r._21+l.z*r._22", metal));
327 genFunc(builder, "vec3", "mul", "mat3", "vec3", _ ->
328 matrixFunctionReturn(builder, "vec3",
329 "l._00*r.x+l._10*r.y+l._20*r.z,l._01*r.x+l._11*r.y+l._21*r.z,l._02*r.x+l._12*r.y+l._22*r.z", metal));
330 genFunc(builder, "mat3", "mul", "mat3", "mat3", _ ->
331 matrixFunctionReturn(builder, "mat3",
332 "l._00*r._00,l._01*r._01,l._02*r._02,l._10*r._10,l._11*r._11,l._12*r._12,l._20*r._20,l._21*r._21,l._22*r._22", metal));
333 genFunc(builder, "mat2", "mul", "mat2", "mat2", _ ->
334 matrixFunctionReturn(builder, "mat2", "l._00*r._00,l._01*r._01,l._10*r._10,l._11*r._11", metal));
335 }
336
337 private static <T extends C99HATKernelBuilder<T>> void matrixFunctionReturn(
338 C99HATKernelBuilder<T> builder, String type, String expression, boolean metal) {
339 builder.returnKeyword().sp();
340 // builder.type(type);
341 builder.paren(_ -> builder.type(type));
342
343 if (type.startsWith("mat")) {
344 builder.brace(_ -> builder.preformatted(expression));
345 } else {
346 builder.paren(_ -> builder.preformatted(expression));
347 }
348 builder.semicolon();
349 }
350
351 private static <T extends C99HATKernelBuilder<T>> void handleRoundInvoke(
352 C99HATKernelBuilder<T> builder, OpHelper.Invoke invoke) {
353 nameAndArgs(invoke, clName(invoke.lookup(), (JavaType) invoke.returnType()) + "_round", builder);
354 }
355
356 public static <T extends C99HATKernelBuilder<T>> void createRoundFunctions(
357 C99HATKernelBuilder<T> builder, boolean metal) {
358 // F32.round uses Math.round(float): ties toward +infinity and an int result converted to float.
359 builder.func(_ -> builder.type("float"), "float_round",
360 _ -> builder.type("float").sp().id("x"),
361 _ -> builder.preformatted("""
362 if (isnan(x)) return 0.0f;
363 if (x >= 2147483648.0f) return 2147483648.0f;
364 if (x <= -2147483648.0f) return -2147483648.0f;
365 float base = floor(x);
366 return (float)((int)base + (x - base >= 0.5f ? 1 : 0));
367 """)).semicolon();
368 List.of("vec2", "vec3", "vec4").forEach(type -> {
369 String lanes = String.join(",", List.of("x", "y", "z", "w")
370 .subList(0, Integer.parseInt(type.substring(3))).stream()
371 .map(lane -> "float_round(v." + lane + ")").toList());
372 builder.func(_ -> builder.type(type), type + "_round",
373 _ -> builder.type(type).sp().id("v"),
374 _ -> matrixFunctionReturn(builder, type, lanes, metal)).semicolon();
375 });
376 }
377
378 public static <T extends C99HATKernelBuilder<T>> void createModFunctions(C99HATKernelBuilder<T> builder) {
379 List.of("float", "vec2", "vec3", "vec4").forEach(type -> {
380 createModFunction(builder, type, "float");
381 if (!type.equals("float")) {
382 createModFunction(builder, type, type);
383 }
384 });
385 }
386
387 private static <T extends C99HATKernelBuilder<T>> void createModFunction(
388 C99HATKernelBuilder<T> builder, String lhs, String rhs) {
389 genFunc(builder, lhs, "mod", lhs, rhs, _ ->
390 builder.returnKeyword().sp().id("l").sp().minus().id("r").sp().mul().sp().id("floor")
391 .paren(_ -> builder.id("l").div().id("r")).semicolon());
392 }
393
394
395 }