1 /*
2 * Copyright (c) 2025, 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 oracle.code.samples;
26
27 import jdk.incubator.code.AbstractOp;
28 import jdk.incubator.code.CodeContext;
29 import jdk.incubator.code.CodeTransformer;
30 import jdk.incubator.code.Reflect;
31 import jdk.incubator.code.Op;
32 import jdk.incubator.code.CodeType;
33 import jdk.incubator.code.Value;
34 import jdk.incubator.code.dialect.core.CoreOp;
35 import jdk.incubator.code.dialect.core.SSA;
36 import jdk.incubator.code.dialect.java.JavaOp;
37 import jdk.incubator.code.extern.ExternalizedOp;
38
39 import java.lang.reflect.Method;
40 import java.util.List;
41 import java.util.Optional;
42 import java.util.stream.Stream;
43
44 /**
45 * Simple example demonstrating how to create new nodes to build a new dialect.
46 * <p>
47 * This example customizes invoke ops to handle them as intrinsics operations.
48 * The example showcase how to search for Java methods with specific signature (in this case
49 * FMA), and replace them with Op that performs the FMA operation.
50 * </p>
51 *
52 * <p>
53 * How to run from the terminal?
54 * <code>
55 * $ java --add-modules jdk.incubator.code -cp target/crsamples-1.0-SNAPSHOT.jar oracle.code.samples.DialectWithInvoke
56 * </code>
57 * </p>
58 */
59 public class DialectWithInvoke {
60
61 private static float intrinsicsFMA(float a, float b, float c) {
62 return Math.fma(a, b, c);
63 }
64
65 @Reflect
66 public static float myFunction(float a, float b, float c) {
67 return intrinsicsFMA(a, b, c);
68 }
69
70 // Custom/Dialect Nodes extends from Op
71 public static class FMAIntrinsicOp extends AbstractOp implements ExternalizedOp.Externalizable { // externalized
72
73 private final CodeType typeDescriptor;
74
75 FMAIntrinsicOp(CodeType typeDescriptor, List<Value> operands) {
76 super(operands);
77 this.typeDescriptor = typeDescriptor;
78 }
79
80 FMAIntrinsicOp(FMAIntrinsicOp that, CodeContext cc) {
81 super(that, cc);
82 this.typeDescriptor = that.resultType();
83 }
84
85 @Override
86 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
87 return new FMAIntrinsicOp(this, copyContext);
88 }
89
90 @Override
91 public CodeType resultType() {
92 return typeDescriptor;
93 }
94
95 @Override
96 public String externalizeOpName() {
97 return "intrinsicsFMA";
98 }
99 }
100
101 private static void customInvoke() {
102
103 Optional<Method> myFunction = Stream.of(DialectWithInvoke.class.getDeclaredMethods())
104 .filter(m -> m.getName().equals("myFunction"))
105 .findFirst();
106 Method m = myFunction.get();
107
108 // Original Code Mode.
109 CoreOp.FuncOp functionModel = Op.ofMethod(m).get();
110 IO.println(functionModel.toText());
111
112 // Transform the code model to search for all InvokeOp and check if the
113 // method name matches with the one we want to replace. We could also check
114 // parameters and their types. For simplication, this example does not check this.
115 CoreOp.FuncOp dialectModel = functionModel.transform((blockBuilder, op) -> {
116 CodeContext context = blockBuilder.context();
117 if (op instanceof JavaOp.InvokeOp invokeOp && invokeOp.invokeReference().name().equals("intrinsicsFMA")) {
118 // The Op is the one we are looking for.
119 // We obtain the input values to this Op and use them to build the new FMA op.
120 List<Value> inputOperands = invokeOp.operands();
121 List<Value> outputOperands = context.getValues(inputOperands);
122
123 // Create new node
124 FMAIntrinsicOp myCustomFunction = new FMAIntrinsicOp(invokeOp.resultType(), outputOperands);
125
126 // Add the new node to the code builder
127 Op.Result outputResult = blockBuilder.add(myCustomFunction);
128
129 // Preserve the location from the original invoke
130 myCustomFunction.setLocation(invokeOp.location());
131
132 // Map input-> new output
133 context.mapValue(invokeOp.result(), outputResult);
134 } else {
135 blockBuilder.add(op);
136 }
137 return blockBuilder;
138 });
139
140 IO.println("Model with new OpNodes for Dialect: ");
141 IO.println(dialectModel.toText());
142
143 CoreOp.FuncOp ssaDialect = SSA.transform(dialectModel);
144 IO.println("Model with new OpNodes for SsaDialect: ");
145 IO.println(ssaDialect.toText());
146 }
147
148 static void main() {
149 IO.println("Testing Dialects in Code-Reflection");
150 customInvoke();
151 }
152 }