1 /*
2 * Copyright (c) 2024, 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
26 package jdk.incubator.code.analysis;
27
28 import jdk.incubator.code.*;
29 import jdk.incubator.code.op.AnfDialect;
30 import jdk.incubator.code.op.CoreOp;
31 import jdk.incubator.code.type.FunctionType;
32 import java.util.*;
33 import java.util.function.Function;
34
35 public class AnfTransformer {
36
37
38 final CoreOp.FuncOp sourceOp;
39 final Map<Block, Function<Body.Builder, AnfDialect.AnfFuncOp>> fBuilders = new HashMap<>();
40 final Body.Builder outerBodyBuilder;
41 final ImmediateDominatorMap idomMap;
42 final Map<Block, Value> funMap = new HashMap<>();
43 final Map<Block, Value> funMap2 = new HashMap<>();
44
45 public AnfTransformer(CoreOp.FuncOp funcOp) {
46 sourceOp = funcOp;
47 outerBodyBuilder = Body.Builder.of(null, FunctionType.functionType(funcOp.body().yieldType()));
48 idomMap = new ImmediateDominatorMap(funcOp.body());
49 }
50
51 public AnfDialect.AnfFuncOp transform() {
52 return transformOuterBody(sourceOp.body());
53 }
54
55 //Outer body corresponds to outermost letrec
56 public AnfDialect.AnfFuncOp transformOuterBody(Body b) {
57 var entry = b.entryBlock();
58
59 var builderEntry = outerBodyBuilder.entryBlock();
60
61 var selfRefP = builderEntry.parameter(((CoreOp.FuncOp) b.parentOp()).invokableType());
62 funMap.put(entry, selfRefP);
63
64 for (Block.Parameter p : entry.parameters()) {
65 var newP = builderEntry.parameter(p.type());
66 builderEntry.context().mapValue(p,newP);
67 }
68
69 var outerLetRecBody = Body.Builder.of(outerBodyBuilder, FunctionType.functionType(b.yieldType(), List.of()), CopyContext.create(builderEntry.context()));
70
71 List<Block> dominatedBlocks = idomMap.idominates(entry);
72 List<AnfDialect.AnfFuncOp> funs = dominatedBlocks.stream().map(block -> transformBlock(block, outerLetRecBody)).toList();
73
74 var res = transformBlock(entry, outerLetRecBody);
75 return res;
76
77 }
78
79 public AnfDialect.AnfFuncOp transformBlock(Block b, Body.Builder bodyBuilder) {
80 if (idomMap.idominates(b).isEmpty()) {
81 return transformLeafBlock(b, bodyBuilder);
82 }
83 return transformDomBlock(b, bodyBuilder);
84 }
85
86 //"Leaf" in this case is a leaf of the dominator tree
87 public AnfDialect.AnfFuncOp transformLeafBlock(Block b, Body.Builder ancestorBodyBuilder) {
88 var blockReturnType = getBlockReturnType(b);
89 var blockFType = FunctionType.functionType(blockReturnType);
90
91 List<TypeElement> synthParamTypes = new ArrayList<>();
92 synthParamTypes.add(blockFType);
93
94 var blockFTypeSynth = FunctionType.functionType(blockReturnType, synthParamTypes);
95
96 Body.Builder newBodyBuilder = Body.Builder.of(ancestorBodyBuilder, blockFTypeSynth, CopyContext.create(ancestorBodyBuilder.entryBlock().context()));
97
98 var selfRefParam = newBodyBuilder.entryBlock().parameters().get(0);
99 funMap.put(b, selfRefParam);
100
101 for (Block.Parameter param : b.parameters()) {
102 var p = newBodyBuilder.entryBlock().parameter(param.type());
103 newBodyBuilder.entryBlock().context().mapValue(param, p);
104 }
105
106 var letBody = Body.Builder.of(newBodyBuilder, FunctionType.functionType(blockReturnType, List.of()), CopyContext.create(newBodyBuilder.entryBlock().context()));
107
108 AnfDialect.AnfLetOp let = transformOps(b, letBody);
109 newBodyBuilder.entryBlock().op(let);
110 return AnfDialect.func(b.toString(), newBodyBuilder);
111 }
112
113 //Non leaf nodes of the dominator tree
114 public AnfDialect.AnfFuncOp transformDomBlock(Block b, Body.Builder ancestorBodyBuilder) {
115 var blockReturnType = getBlockReturnType(b);
116 var blockFType = FunctionType.functionType(blockReturnType);
117
118 List<TypeElement> synthParamTypes = new ArrayList<>();
119 synthParamTypes.add(blockFType);
120
121 var blockFTypeSynth = FunctionType.functionType(blockReturnType, synthParamTypes);
122
123 //Function body contains letrec and its bodies
124 Body.Builder funcBodyBuilder = Body.Builder.of(ancestorBodyBuilder, blockFTypeSynth, CopyContext.create(ancestorBodyBuilder.entryBlock().context()));
125
126 //Self param
127 var selfRefParam = funcBodyBuilder.entryBlock().parameters().get(0);
128 funMap.put(b, selfRefParam);
129
130 for (Block.Parameter param : b.parameters()) {
131 var p = funcBodyBuilder.entryBlock().parameter(param.type());
132 funcBodyBuilder.entryBlock().context().mapValue(param, p);
133 }
134
135 //letrec inner body
136 Body.Builder letrecBody = Body.Builder.of(funcBodyBuilder, FunctionType.functionType(blockReturnType, List.of()), CopyContext.create(funcBodyBuilder.entryBlock().context()));
137
138 List<Block> dominates = idomMap.idominates(b);
139 for (Block dblock : dominates) {
140 var res = transformDomBlock(dblock, letrecBody);
141 var fval = letrecBody.entryBlock().op(res);
142 funMap2.put(dblock, fval);
143 }
144
145 var letBody = Body.Builder.of(letrecBody, letrecBody.bodyType(), CopyContext.create(letrecBody.entryBlock().context()));
146 transformBlockOps(b, letBody.entryBlock());
147 var let = AnfDialect.let(letBody);
148
149 letrecBody.entryBlock().op(let);
150
151 var letrec = AnfDialect.letrec(letrecBody);
152 funcBodyBuilder.entryBlock().op(letrec);
153 return AnfDialect.func(b.toString(), funcBodyBuilder);
154
155 }
156
157 private TypeElement getBlockReturnType(Block b) {
158 var op = b.ops().getLast();
159 if (op instanceof Op.Terminating) {
160 List<Block.Reference> destBlocks = new ArrayList<>();
161 if (op instanceof CoreOp.ReturnOp ro) {
162 return ro.returnValue().type();
163 } else if (op instanceof CoreOp.YieldOp yo) {
164 return yo.yieldValue().type();
165 } else if (op instanceof CoreOp.BranchOp bop) {
166 destBlocks.addAll(bop.successors());
167 } else if (op instanceof CoreOp.ConditionalBranchOp cbop) {
168 destBlocks.addAll(cbop.successors());
169 }
170 //Traverse until we find a yield or return type, TODO: not going to try to unify types
171
172 Set<Block> visitedBlocks = new HashSet<>();
173 visitedBlocks.add(b);
174
175 while (!destBlocks.isEmpty()) {
176 var block = destBlocks.removeFirst().targetBlock();
177 if (visitedBlocks.contains(block)) {
178 continue;
179 }
180
181 //Discovered a terminator with a return value, use its type
182 if (block.successors().isEmpty()) {
183 var o = block.ops().getLast();
184 if (o instanceof CoreOp.ReturnOp ro) {
185 return ro.returnValue().type();
186 } else if (o instanceof CoreOp.YieldOp yo) {
187 return yo.yieldValue().type();
188 } else {
189 throw new UnsupportedOperationException("Unsupported terminator encountered: " + o.opName());
190 }
191 } else {
192 visitedBlocks.add(block);
193 var newDests = block.successors().stream().filter((s) -> !visitedBlocks.contains(s.targetBlock())).toList();
194 destBlocks.addAll(newDests);
195 }
196 }
197
198 }
199
200 throw new RuntimeException("Encountered Block with no return " + op.opName());
201 }
202
203 private Block.Builder transformEndOp(Block.Builder b, Op op) {
204 if (op instanceof Op.Terminating t) {
205 switch (t) {
206 case CoreOp.ConditionalBranchOp c -> {
207 var tbranch_args = c.trueBranch().arguments();
208 tbranch_args = tbranch_args.stream().map(b.context()::getValue).toList();
209 var fbranch_args = c.falseBranch().arguments();
210 fbranch_args = fbranch_args.stream().map(b.context()::getValue).toList();
211
212 List<Value> trueArgs = new ArrayList<>();
213 trueArgs.addAll(tbranch_args);
214
215 List<Value> falseArgs = new ArrayList<>();
216 falseArgs.addAll(fbranch_args);
217
218
219 var ifExp = AnfDialect.if_(b.parentBody(),
220 getBlockReturnType(c.trueBranch().targetBlock()),
221 b.context().getValue(c.predicate()))
222 .if_((bodyBuilder) -> bindFunApp(bodyBuilder, trueArgs, c.trueBranch().targetBlock()))
223 .else_((bodyBuilder) -> bindFunApp(bodyBuilder, falseArgs, c.falseBranch().targetBlock()));
224
225 b.op(ifExp);
226
227 return b;
228 }
229 case CoreOp.BranchOp br -> {
230 var args = br.branch().arguments();
231 args = args.stream().map(b.context()::getValue).toList();
232
233 List<Value> funcArgs = new ArrayList<>();
234 funcArgs.addAll(args);
235 bindFunApp(b, funcArgs, br.branch().targetBlock());
236
237 return b;
238 }
239 case CoreOp.ReturnOp ro -> {
240 var rval = b.context().getValue(ro.returnValue());
241 b.op(CoreOp._yield(rval));
242 return b;
243 }
244 case CoreOp.YieldOp y -> {
245 var rval = b.context().getValue(y.yieldValue());
246 b.op(CoreOp._yield(rval));
247 return b;
248 }
249 default -> {
250 throw new UnsupportedOperationException("Unsupported terminating op encountered: " + op);
251 }
252 }
253 } else {
254 b.op(op);
255 return b;
256 }
257 }
258
259
260 private void bindFunApp(Block.Builder b, List<Value> args, Block target) {
261
262 List<Value> synthArgs = new ArrayList<>();
263 synthArgs.addAll(args);
264 synthArgs.addFirst(funMap.get(target));
265 try {
266 b.op(AnfDialect.apply(synthArgs));
267 return;
268 } catch (IllegalStateException e) {}
269
270 synthArgs.removeFirst();
271 synthArgs.addFirst(funMap2.get(target));
272
273 try {
274 b.op(AnfDialect.apply(synthArgs));
275 } catch (IllegalStateException e) {
276 throw new IllegalStateException("No valid mapping to FuncOp for apply");
277 }
278
279 }
280
281
282 public AnfDialect.AnfLetOp transformOps(Block b, Body.Builder bodyBuilder) {
283 Block.Builder blockb = bodyBuilder.entryBlock();
284 return transformOps(b, blockb);
285 }
286
287 public AnfDialect.AnfLetOp transformOps(Block b, Block.Builder blockBuilder) {
288 transformBlockOps(b, blockBuilder);
289 return AnfDialect.let(blockBuilder.parentBody());
290 }
291
292 public void transformBlockOps(Block b, Block.Builder blockBuilder) {
293 for (var op : b.ops()) {
294 transformEndOp(blockBuilder, op);
295 }
296 }
297
298 static class ImmediateDominatorMap {
299
300 private final Map <Block, List<Block>> dominatesMap;
301 private final Map <Block, Block> dominatorsMap;
302
303 public ImmediateDominatorMap(Body b) {
304 dominatorsMap = b.immediateDominators();
305 dominatesMap = new HashMap<>();
306
307 //Reverse the idom relation
308 b.immediateDominators().forEach((dominated, dominator) -> {
309 if (!dominated.equals(dominator)) {
310 dominatesMap.compute(dominator, (k, v) -> {
311 if (v == null) {
312 var newList = new ArrayList<Block>();
313 newList.add(dominated);
314 return newList;
315 } else {
316 v.add(dominated);
317 return v;
318 }
319 });
320 }
321 });
322
323 }
324
325 //Looks "down" the dominator tree toward leaves
326 public List<Block> idominates(Block b) {
327 return dominatesMap.getOrDefault(b, List.of());
328 }
329
330 //Looks "up" the dominator tree toward start node
331 public Block idominatedBy(Block b) {
332 return dominatorsMap.get(b);
333 }
334 }
335 }