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