1 /*
2 * Copyright (c) 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.phases;
26
27 import hat.HATMath;
28 import hat.types.S16ImplOfF16;
29 import jdk.incubator.code.Op;
30 import jdk.incubator.code.dialect.core.CoreOp;
31 import jdk.incubator.code.dialect.java.ClassType;
32 import optkl.OpHelper;
33 import optkl.Trxfmr;
34 import optkl.VarTable;
35
36 import java.lang.invoke.MethodHandles;
37 import java.util.HashMap;
38 import java.util.Map;
39 import java.util.Objects;
40
41 public record HATMathLibPhase() implements HATPhase {
42 @Override
43 public CoreOp.FuncOp transform(MethodHandles.Lookup lookup,CoreOp.FuncOp funcOp, VarTable varTable) {
44 Map<Op, Class<? extends S16ImplOfF16>> setTypeMap = new HashMap<>();
45 OpHelper.Invoke.stream(lookup, funcOp)
46 .filter(invoke -> !invoke.returnsVoid() && invoke.returnsClassType() && invoke.refIs(HATMath.class))
47 .forEach(invoke ->
48 // This detects a HATMathLib is stored either in a VarOp or a VarStoreOp
49 invoke.op().result().uses().stream()
50
51 .filter(result -> (result.op() instanceof CoreOp.VarOp) || (result.op() instanceof CoreOp.VarAccessOp.VarStoreOp))
52 .findFirst()
53 .ifPresent(result -> {
54 // Special attention to HATTypes. This is some metadata associated with the invoke
55 // to pass to the cogen to do further processing (e.g., build a new type or typecast).
56 // An alternative is to insert a `stub`, or a code snippet that insert new nodes in the
57 // IR
58 if (S16ImplOfF16.codeTypeToFloatClassOrNull(invoke,(ClassType)invoke.returnType()) instanceof Class<? extends S16ImplOfF16> reducedFloatType) {
59 setTypeMap.put(result.op(), reducedFloatType);
60 setTypeMap.put(invoke.op(), reducedFloatType);
61 }
62 }));
63 return Trxfmr.of(lookup, funcOp).transform(setTypeMap::containsKey, (blockBuilder, op) -> {
64 if (Objects.requireNonNull(op) instanceof CoreOp.VarOp varOp) {
65 if (setTypeMap.get(varOp) == null) {
66 // this varOp is not a special type (e.g., float16), then we insert the varOp into the new tree
67 blockBuilder.add(varOp);
68 } else {
69 // Add the special type as a VarOp
70 HATFP16Phase.copyVarOpWithUpdateVarTable(funcOp.funcName(), varOp, blockBuilder, varTable);
71 // If we add HATMath ops for other special types in HAT (e.g., Vectors),
72 // we may need to also add the new <X>HATVarOps here as well.
73 }
74 } else {
75 blockBuilder.add(op);
76 }
77 return blockBuilder;
78 }, varTable).funcOp();
79 }
80 }