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