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