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 hat.dialect;
 26 
 27 import jdk.incubator.code.Op;
 28 import jdk.incubator.code.TypeElement;
 29 import jdk.incubator.code.Value;
 30 import jdk.incubator.code.dialect.core.CoreOp;
 31 import jdk.incubator.code.dialect.java.ClassType;
 32 import jdk.incubator.code.dialect.java.JavaOp;
 33 import jdk.incubator.code.dialect.java.JavaType;
 34 import optkl.OpHelper;
 35 
 36 import java.lang.invoke.MethodHandles;
 37 import java.lang.reflect.Method;
 38 import java.util.Optional;
 39 import java.util.concurrent.atomic.AtomicInteger;
 40 import java.util.concurrent.atomic.AtomicReference;
 41 import java.util.stream.Stream;
 42 
 43 public class HATPhaseUtils {
 44 
 45     public record VectorMetaData(TypeElement vectorTypeElement, int lanes) {
 46     }
 47 
 48     public static VectorMetaData getVectorTypeInfo(MethodHandles.Lookup lookup, JavaOp.InvokeOp invokeOp, int param) {
 49         Value varValue = invokeOp.operands().get(param);
 50         if (varValue instanceof Op.Result r && r.op() instanceof CoreOp.VarAccessOp.VarLoadOp varLoadOp) {
 51             return getVectorTypeInfoWithCodeReflection(lookup,varLoadOp.resultType());
 52         }
 53         return null;
 54     }
 55 
 56     public static VectorMetaData getVectorTypeInfo(MethodHandles.Lookup lookup,JavaOp.InvokeOp invokeOp) {
 57         return getVectorTypeInfoWithCodeReflection(lookup,invokeOp.resultType());
 58     }
 59     public static TypeElement getVectorElementType(String primitive) {
 60         return switch (primitive) {
 61             case "float" -> JavaType.FLOAT;
 62             case "double" -> JavaType.DOUBLE;
 63             case "int" -> JavaType.INT;
 64             case "long" -> JavaType.LONG;
 65             case "short" -> JavaType.SHORT;
 66             case "byte" -> JavaType.BYTE;
 67             case "char" -> JavaType.CHAR;
 68             case "boolean" -> JavaType.BOOLEAN;
 69             default -> null;
 70         };
 71     }
 72 
 73     /**
 74      * This method inspects the Vector Type Methods to obtain two methods for code-model:
 75      * 1) Method `type` to obtain the primitive base type of the vector type.
 76      * 2) Method `width` to obtain the number of lanes.
 77      *
 78      * @param typeElement
 79      *  {@link TypeElement}
 80      * @return
 81      * {@link VectorMetaData}
 82      */
 83     public static VectorMetaData getVectorTypeInfoWithCodeReflection(MethodHandles.Lookup lookup,TypeElement typeElement) {
 84         Class<?> clazz = (Class<?>) OpHelper.classTypeToTypeOrThrow(lookup, (ClassType) typeElement);
 85         CoreOp.FuncOp codeModelType = buildCodeModelFor(clazz, "type");
 86         AtomicReference<TypeElement> vectorElement = new AtomicReference<>();
 87         codeModelType.elements().forEach(codeElement -> {
 88             if (codeElement instanceof CoreOp.ReturnOp returnOp) {
 89                 Value v = returnOp.operands().getFirst();
 90                 if (v instanceof Op.Result r && r.op() instanceof JavaOp.FieldAccessOp.FieldLoadOp fieldLoadOp) {
 91                     String primitiveTypeName = fieldLoadOp.fieldDescriptor().name();
 92                     vectorElement.set(getVectorElementType(primitiveTypeName.toLowerCase()));
 93                 }
 94             }
 95         });
 96 
 97         AtomicInteger lanes = new AtomicInteger(1);
 98         CoreOp.FuncOp codeModelWidth = buildCodeModelFor(clazz, "width");
 99         codeModelWidth.elements().forEach(codeElement -> {
100             if (codeElement instanceof CoreOp.ReturnOp returnOp) {
101                 Value v = returnOp.operands().getFirst();
102                 if (v instanceof Op.Result r && r.op() instanceof CoreOp.ConstantOp constantOp) {
103                     lanes.set((Integer) constantOp.value());
104                 }
105             }
106         });
107         return new VectorMetaData(vectorElement.get(), lanes.get());
108     }
109 
110 
111     private static CoreOp.FuncOp buildCodeModelFor(Class<?> klass, String methodName) {
112         Optional<Method> methodFunction = Stream.of(klass.getMethods())
113                 .filter(m -> m.getName().equals(methodName))
114                 .findFirst();
115         return Op.ofMethod(methodFunction.get()).get();
116     }
117 
118 }