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.phases;
26
27 import hat.dialect.HATPtrOp;
28 import jdk.incubator.code.CodeType;
29 import jdk.incubator.code.dialect.java.ArrayType;
30 import jdk.incubator.code.dialect.java.ClassType;
31 import jdk.incubator.code.dialect.java.JavaOp;
32 import optkl.IfaceValue;
33 import optkl.OpHelper;
34 import optkl.Trxfmr;
35 import jdk.incubator.code.Op;
36 import jdk.incubator.code.Value;
37 import jdk.incubator.code.dialect.core.CoreOp;
38 import optkl.util.ops.VarLikeOp;
39
40 import java.lang.invoke.MethodHandles;
41 import java.util.ArrayList;
42 import java.util.HashMap;
43 import java.util.HashSet;
44 import java.util.List;
45 import java.util.Map;
46 import java.util.Set;
47
48 import static hat.phases.HATPhaseUtils.findOpInResultFromFirstOperandsOrNull;
49 import static optkl.OpHelper.Invoke;
50 import static optkl.OpHelper.Invoke.invoke;
51 import static optkl.OpHelper.classTypeToTypeOrThrow;
52 import static optkl.OpHelper.copyLocation;
53 import static optkl.OpHelper.firstOperandOrThrow;
54 import static optkl.OpHelper.opFromFirstOperandOrNull;
55 import static optkl.OpHelper.opFromFirstOperandOrThrow;
56 import static optkl.OpHelper.resultFromFirstOperandOrNull;
57 import static optkl.OpHelper.resultFromFirstOperandOrThrow;
58 import static optkl.OpHelper.resultFromOperandN;
59
60 public record HATArrayViewPhase() implements HATPhase {
61 public static boolean isVectorOp(MethodHandles.Lookup lookup, Op op) {
62 if (!op.operands().isEmpty()) {
63 CodeType type = switch (op) {
64 case JavaOp.ArrayAccessOp.ArrayLoadOp load -> load.resultType();
65 case JavaOp.ArrayAccessOp.ArrayStoreOp store -> store.operands().getLast().type();
66 default -> OpHelper.firstOperandOrThrow(op).type();
67 };
68 if (type instanceof ArrayType at) {
69 type = at.componentType();
70 }
71 if (type instanceof ClassType ct) {
72 try {
73 return IfaceValue.Vector.class.isAssignableFrom((Class<?>) ct.resolve(lookup));
74 } catch (ReflectiveOperationException e) {
75 throw new IllegalStateException(e);
76 }
77 }
78 }
79 return false;
80 }
81
82 public static boolean isVectorBinaryOp(MethodHandles.Lookup lookup, OpHelper.Invoke invoke) {
83 return isVectorOp(lookup, invoke.op()) && invoke.nameMatchesRegex("(add|sub|mul|div)");
84 }
85
86 public static boolean isBufferArray(Op op) {
87 JavaOp.InvokeOp iop = (JavaOp.InvokeOp) findOpInResultFromFirstOperandsOrNull(op, JavaOp.InvokeOp.class);
88 return iop != null && iop.invokeReference().name().toLowerCase().contains("arrayview"); // we need a better way
89 }
90
91 public static boolean isBufferInitialize(Op op) {
92 // first check if the return is an array type
93 if (op instanceof CoreOp.VarOp vop && vop.varValueType() instanceof ArrayType
94 || op instanceof JavaOp.ArrayAccessOp
95 || op.resultType() instanceof ArrayType) return isBufferArray(op);
96 return false;
97 }
98
99 public static boolean isLocalSharedOrPrivate(Op op) {
100 JavaOp.InvokeOp iop = (JavaOp.InvokeOp) findOpInResultFromFirstOperandsOrNull(op, JavaOp.InvokeOp.class);
101 return iop != null
102 && (iop.invokeReference().name().toLowerCase().contains("shared")
103 || iop.invokeReference().name().toLowerCase().contains("local")
104 || iop.invokeReference().name().toLowerCase().contains("private")
105 );
106 }
107
108 static HATArrayViewPhase.ArrayAccessInfo arrayAccessInfo(Value value, Map<Op.Result, Op.Result> replaced) {
109 return expressionGraph(value).getInfo(replaced);
110 }
111
112 static HATArrayViewPhase.Node<Value> expressionGraph(Value value) {
113 return expressionGraph(new HashMap<>(), value);
114 }
115
116 static HATArrayViewPhase.Node<Value> expressionGraph(Map<Value, HATArrayViewPhase.Node<Value>> visited, Value value) {
117 // If value has already been visited return its node
118 if (visited.containsKey(value)) {
119 return visited.get(value);
120 }
121
122 // Find the expression graphs for each operand
123 List<HATArrayViewPhase.Node<Value>> edges = new ArrayList<>();
124
125 // looks like
126 for (Value operand : value.dependsOn()) {
127 if (operand instanceof Op.Result res &&
128 res.op() instanceof JavaOp.InvokeOp iop
129 && iop.invokeReference().name().toLowerCase().contains("arrayview")) { // We need to find a better way
130 continue;
131 }
132 edges.add(expressionGraph(operand));
133 }
134 HATArrayViewPhase.Node<Value> node = new HATArrayViewPhase.Node<>(value, edges);
135 visited.put(value, node);
136 return node;
137 }
138
139 @Override
140 public CoreOp.FuncOp transform(MethodHandles.Lookup lookup, CoreOp.FuncOp funcOp, VarTable varTable) {
141 if (Invoke.stream(lookup, funcOp).noneMatch(
142 invoke -> isBufferArray(invoke.op())
143 )) return funcOp;
144
145 funcOp = applyArrayView(lookup, funcOp);
146
147 if (funcOp.elements().filter(e -> e instanceof CoreOp.VarOp).anyMatch(
148 e -> isVectorOp(lookup, ((CoreOp.VarOp) e))
149 )) funcOp = applyVectorView(lookup, funcOp, varTable);
150 return funcOp;
151 }
152
153 public CoreOp.FuncOp applyVectorView(MethodHandles.Lookup lookup, CoreOp.FuncOp funcOp, VarTable varTable) {
154 return Trxfmr.of(lookup, funcOp).transform((blockBuilder, op) -> {
155 switch (op) {
156 case JavaOp.InvokeOp iOp when invoke(lookup, iOp) instanceof Invoke invoke && isVectorBinaryOp(invoke.lookup(), invoke) ->
157 blockBuilder.add(op);
158 case CoreOp.VarOp varOp when isVectorOp(lookup, varOp) -> {
159 Op.Result op1 = blockBuilder.add(varOp);
160 String functionName = funcOp.funcName();
161 varTable.addIfNeededOrThrow(functionName, op1.op(), VarTable.HATOpAttribute.VECTOR);
162 return blockBuilder;
163 }
164 case JavaOp.ArrayAccessOp.ArrayLoadOp arrayLoadOp -> {
165 if (isVectorOp(lookup, arrayLoadOp)) {
166 blockBuilder.add(op);
167 }
168 return blockBuilder;
169 }
170 case JavaOp.ArrayAccessOp.ArrayStoreOp arrayStoreOp -> {
171 if (isVectorOp(lookup, arrayStoreOp)) {
172 blockBuilder.add(op);
173 }
174 return blockBuilder;
175 }
176 default -> {
177 }
178 }
179 blockBuilder.add(op);
180 return blockBuilder;
181 }).funcOp();
182 }
183
184 public CoreOp.FuncOp applyArrayView(MethodHandles.Lookup lookup, CoreOp.FuncOp funcOp) {
185 Map<Op.Result, Op.Result> replaced = new HashMap<>(); // maps a result to the result it should be replaced by
186 Map<Op, CoreOp.VarAccessOp.VarLoadOp> bufferVarLoads = new HashMap<>();
187
188 return Trxfmr.of(lookup, funcOp).transform((blockBuilder, op) -> {
189 var context = blockBuilder.context();
190 switch (op) {
191 case JavaOp.InvokeOp invokeOp when invoke(lookup, invokeOp) instanceof Invoke invoke && isBufferArray(invoke.op()) -> {
192 Op.Result result = invoke.resultFromFirstOperandOrNull();
193 replaced.put(invoke.returnResult(), result);
194 // map buffer VarOp to its corresponding VarLoadOp
195 bufferVarLoads.put((opFromFirstOperandOrNull(result.op())), (CoreOp.VarAccessOp.VarLoadOp) result.op());
196 return blockBuilder;
197 }
198 case CoreOp.VarOp varOp when isBufferInitialize(varOp) -> {
199 Op bufferLoad = replaced.get(resultFromFirstOperandOrThrow(varOp)).op(); // gets VarLoadOp associated w/ og buffer
200 replaced.put(varOp.result(), resultFromFirstOperandOrNull(bufferLoad)); // gets VarOp associated w/ og buffer
201 return blockBuilder;
202 }
203 case CoreOp.VarAccessOp.VarLoadOp varLoadOp when (isBufferInitialize(varLoadOp)) -> {
204 Op.Result r = resultFromFirstOperandOrThrow(varLoadOp);
205 Op.Result replacement;
206 if (r.op() instanceof CoreOp.VarOp) { // if this is the VarLoadOp after the .arrayView() InvokeOp
207 replacement = (isLocalSharedOrPrivate(varLoadOp)) ?
208 resultFromFirstOperandOrNull(opFromFirstOperandOrThrow(r.op())) :
209 bufferVarLoads.get(replaced.get(r).op()).result();
210 } else { // if this is a VarLoadOp loading the buffer
211 CoreOp.VarAccessOp.VarLoadOp newVarLoad = CoreOp.varLoad(blockBuilder.context().getValue(replaced.get(r)));
212 replacement = blockBuilder.add(copyLocation(varLoadOp, newVarLoad));
213 context.mapValue(varLoadOp.result(), replacement);
214 }
215 replaced.put(varLoadOp.result(), replacement);
216 return blockBuilder;
217 }
218 case JavaOp.ArrayAccessOp.ArrayLoadOp arrayLoadOp when isBufferArray(arrayLoadOp) -> {
219 Op replacementOp;
220 if (isVectorOp(lookup, arrayLoadOp)) {
221 replacementOp = JavaOp.arrayLoadOp(
222 context.getValue(replaced.get((Op.Result) arrayLoadOp.operands().getFirst())),
223 context.getValue(arrayLoadOp.operands().getLast()),
224 arrayLoadOp.resultType()
225 );
226 } else if (((ArrayType) firstOperandOrThrow(op).type()).dimensions() == 1) {
227 var arrayAccessInfo = arrayAccessInfo(op.result(), replaced);
228 var operands = arrayAccessInfo.bufferAndIndicesAsValues();
229 replacementOp = new HATPtrOp.HATPtrLoadOp(
230 arrayAccessInfo.bufferName(),
231 arrayLoadOp.resultType(),
232 (Class<?>) classTypeToTypeOrThrow(lookup, (ClassType) arrayAccessInfo.buffer().type()),
233 context.getValues(operands)
234 );
235 } else { // we only use the last array load
236 return blockBuilder;
237 }
238 context.mapValue(arrayLoadOp.result(), blockBuilder.add(copyLocation(arrayLoadOp, replacementOp)));
239 return blockBuilder;
240 }
241 case JavaOp.ArrayAccessOp.ArrayStoreOp arrayStoreOp when isBufferArray(arrayStoreOp) -> {
242 Op replacementOp;
243 if (isVectorOp(lookup, arrayStoreOp)) {
244 replacementOp = JavaOp.arrayStoreOp(
245 context.getValue(replaced.get((Op.Result) arrayStoreOp.operands().getFirst())),
246 context.getValue(arrayStoreOp.operands().get(1)),
247 context.getValue(arrayStoreOp.operands().getLast())
248 );
249 } else if (((ArrayType) firstOperandOrThrow(op).type()).dimensions() == 1) { // we only use the last array load
250 var arrayAccessInfo = arrayAccessInfo(op.result(), replaced);
251 var operands = arrayAccessInfo.bufferAndIndicesAsValues();
252 operands.add(arrayStoreOp.operands().getLast());
253 replacementOp = new HATPtrOp.HATPtrStoreOp(
254 arrayAccessInfo.bufferName(),
255 arrayStoreOp.resultType(),
256 (Class<?>) classTypeToTypeOrThrow(lookup, (ClassType) arrayAccessInfo.buffer().type()),
257 context.getValues(operands)
258 );
259 } else {
260 return blockBuilder;
261 }
262 context.mapValue(arrayStoreOp.result(), blockBuilder.add(copyLocation(arrayStoreOp, replacementOp)));
263 return blockBuilder;
264 }
265 case JavaOp.ArrayLengthOp arrayLengthOp when
266 isBufferArray(arrayLengthOp) && resultFromFirstOperandOrThrow(arrayLengthOp) != null -> {
267 var arrayAccessInfo = arrayAccessInfo(op.result(), replaced);
268 var hatPtrLengthOp = new HATPtrOp.HATPtrLengthOp(
269 arrayAccessInfo.bufferName(),
270 arrayLengthOp.resultType(),
271 (Class<?>) OpHelper.classTypeToTypeOrThrow(lookup, (ClassType) arrayAccessInfo.buffer().type()),
272 context.getValues(List.of(arrayAccessInfo.buffer()))
273 );
274 context.mapValue(arrayLengthOp.result(), blockBuilder.add(copyLocation(arrayLengthOp, hatPtrLengthOp)));
275 return blockBuilder;
276 }
277 default -> {
278 }
279 }
280 blockBuilder.add(op);
281 return blockBuilder;
282 }).funcOp();
283 }
284
285 record ArrayAccessInfo(Op.Result buffer, String bufferName, List<Op.Result> indices) {
286 public List<Value> bufferAndIndicesAsValues() {
287 List<Value> operands = new ArrayList<>(List.of(buffer));
288 operands.addAll(indices);
289 return operands;
290 }
291 }
292
293 record Node<T>(T value, List<Node<T>> edges) {
294 ArrayAccessInfo getInfo(Map<Op.Result, Op.Result> replaced) {
295 List<Node<T>> nodeList = new ArrayList<>(List.of(this));
296 Set<Node<T>> handled = new HashSet<>();
297 Op.Result buffer = null;
298 List<Op.Result> indices = new ArrayList<>();
299 while (!nodeList.isEmpty()) {
300 Node<T> node = nodeList.removeFirst();
301 handled.add(node);
302 if (node.value instanceof Op.Result res &&
303 (res.op() instanceof JavaOp.ArrayAccessOp || res.op() instanceof JavaOp.ArrayLengthOp)) {
304 buffer = res;
305 // idx location differs between ArrayAccessOp and ArrayLengthOp
306 indices.addFirst(res.op() instanceof JavaOp.ArrayAccessOp
307 ? resultFromOperandN(res.op(), 1)
308 : resultFromOperandN(res.op(), 0)
309 );
310 }
311 if (!node.edges().isEmpty()) {
312 Node<T> next = node.edges().getFirst(); // we only traverse through the index-related ops
313 if (!handled.contains(next)) {
314 nodeList.add(next);
315 }
316 }
317 }
318 if (buffer != null) {
319 buffer = replaced.get(resultFromFirstOperandOrNull(buffer.op()));
320 String bufferName = hatPtrName(opFromFirstOperandOrNull(buffer.op()));
321 return new ArrayAccessInfo(buffer, bufferName, indices);
322 } else {
323 return null;
324 }
325 }
326 }
327
328 public static String hatPtrName(Op op) {
329 return switch (op) {
330 case CoreOp.VarOp varOp -> varOp.varName();
331 case VarLikeOp varLikeOp -> varLikeOp.varName();
332 case null, default -> "";
333 };
334 }
335 }