lib/choir/src/passes/cse/pass.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const alloc_observe = @import("alloc_observe");
  3 const ir = @import("../../core/root.zig");
  4 const passes = @import("../root.zig");
  5 const pass_mod = passes.pass;
  6 const registry_mod = passes.pipeline;
  7 const control_flow = passes.control_flow;
  8 const cse = @import("root.zig");
  9 
 10 const Pass = pass_mod.Pass;
 11 const PassContext = pass_mod.PassContext;
 12 const PassResult = pass_mod.PassResult;
 13 const Workspace = cse.Workspace;
 14 const Table = cse.Table;
 15 
 16 pub const common_subexpression_elimination_pass_name = "choir-cse";
 17 pub const common_subexpression_elimination_pass_description =
 18     "Common subexpression elimination for pure operations";
 19 
 20 pub fn createCommonSubexpressionEliminationPass() Pass {
 21     return .{
 22         .name = common_subexpression_elimination_pass_name,
 23         .description = common_subexpression_elimination_pass_description,
 24         .run_fn = run,
 25         .mutation_scope = .isolated,
 26         .rerun_policy = .skip_if_unchanged,
 27     };
 28 }
 29 
 30 pub const common_subexpression_elimination_pass_registration = registry_mod.PassRegistration{
 31     .name = common_subexpression_elimination_pass_name,
 32     .description = common_subexpression_elimination_pass_description,
 33     .pass = createCommonSubexpressionEliminationPass(),
 34 };
 35 
 36 pub fn run(ctx: *PassContext) PassResult {
 37     var stack_buffer: [cse.inline_workspace_bytes]u8 = undefined;
 38     var stack_fallback = alloc_observe.buffer.First.init(&stack_buffer, ctx.allocator);
 39     const workspace_allocator = stack_fallback.allocator();
 40     var workspace = Workspace.initForRoot(workspace_allocator, ctx.op) catch
 41         return .failure;
 42     defer workspace.deinit(workspace_allocator);
 43     const dominance: ?*const control_flow.DominanceAnalysis = if (workspace.requiresDominance())
 44         control_flow.getDominanceAnalysis(ctx, ctx.op) catch return .failure
 45     else
 46         null;
 47     workspace.activate() catch return .failure;
 48 
 49     var modified = false;
 50     var replacements: u64 = 0;
 51     cseOnOp(
 52         ctx.op,
 53         dominance,
 54         &workspace,
 55         &modified,
 56         &replacements,
 57     );
 58     std.debug.assert(workspace.insertion_count == 0);
 59     std.debug.assert(workspace.active_scope_count == 0);
 60 
 61     if (replacements != 0) {
 62         ctx.addStatistic("replacements", "redundant operations replaced", replacements);
 63     }
 64 
 65     if (modified) {
 66         ctx.preserveAnalysisSet(control_flow.analysis_ids);
 67         ctx.markModified();
 68     } else {
 69         ctx.preserveAllAnalyses();
 70     }
 71     return .success;
 72 }
 73 
 74 fn cseOnOp(
 75     op: *ir.Operation,
 76     dominance: ?*const control_flow.DominanceAnalysis,
 77     workspace: *Workspace,
 78     modified: *bool,
 79     replacements: *u64,
 80 ) void {
 81     for (op.regions.items) |*region| {
 82         const table = workspace.acquire();
 83         defer workspace.release(table);
 84 
 85         var block_iter = region.getBlocks();
 86         while (block_iter.next()) |block| {
 87             var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
 88             while (current) |current_op| {
 89                 const next = current_op.next_op;
 90 
 91                 if (current_op.regions.items.len > 0) {
 92                     cseOnOp(
 93                         current_op,
 94                         dominance,
 95                         workspace,
 96                         modified,
 97                         replacements,
 98                     );
 99                 }
100 
101                 cseOperation(
102                     current_op,
103                     dominance,
104                     table,
105                     workspace.effect_records,
106                     modified,
107                     replacements,
108                 );
109                 current = next;
110             }
111         }
112     }
113 }
114 
115 fn findDominatingCseCandidate(
116     table: Table,
117     key: u64,
118     op: *ir.Operation,
119     options: ir.OperationEquivalence.Options,
120     dominance: ?*const control_flow.DominanceAnalysis,
121 ) ?*ir.Operation {
122     var candidates = table.candidates(key);
123     while (candidates.next()) |candidate| {
124         if (!cseCandidateDominates(candidate, op, dominance)) continue;
125         if (ir.OperationEquivalence.isEquivalentTo(
126             candidate,
127             op,
128             options,
129         )) {
130             return candidate;
131         }
132     }
133     return null;
134 }
135 
136 fn cseOperation(
137     current_op: *ir.Operation,
138     dominance: ?*const control_flow.DominanceAnalysis,
139     table: Table,
140     effect_records: []ir.interfaces.effects.Fact,
141     modified: *bool,
142     replacements: *u64,
143 ) void {
144     const traits = current_op.getTraits();
145     if (isCseCandidate(current_op, traits, effect_records)) {
146         const commutative = isCommutativeForCse(current_op, traits);
147         const options = cseEquivalenceOptions(commutative);
148         const key = ir.OperationEquivalence.computeHash(current_op, options);
149         const existing = findDominatingCseCandidate(
150             table,
151             key,
152             current_op,
153             options,
154             dominance,
155         );
156         if (existing) |existing_op| {
157             replaceOperation(current_op, existing_op);
158             modified.* = true;
159             replacements.* +|= 1;
160         } else {
161             table.add(key, current_op) catch
162                 @panic("CSE exact candidate capacity exhausted after activation");
163         }
164     }
165 }
166 
167 fn replaceOperation(current_op: *ir.Operation, existing_op: *ir.Operation) void {
168     std.debug.assert(current_op.getNumResults() == existing_op.getNumResults());
169     for (current_op.results.items, existing_op.results.items) |*current_result, *existing_result| {
170         current_result.replaceAllUsesWith(existing_result);
171     }
172     current_op.erase();
173 }
174 
175 fn cseCandidateDominates(
176     candidate: *ir.Operation,
177     op: *ir.Operation,
178     dominance: ?*const control_flow.DominanceAnalysis,
179 ) bool {
180     const candidate_block = candidate.parent_block orelse return false;
181     const op_block = op.parent_block orelse return false;
182     const candidate_region = candidate_block.parent orelse return false;
183     if (op_block.parent != candidate_region) return false;
184     if (candidate_block == op_block) return candidate.isBeforeInBlock(op);
185     const analysis = dominance orelse
186         @panic("cross-block CSE candidate admitted without dominance");
187     return analysis.dominatesBlock(candidate_block, op_block);
188 }
189 
190 fn isCseCandidate(
191     op: *ir.Operation,
192     traits: ir.OperationTraits,
193     storage: []ir.interfaces.effects.Fact,
194 ) bool {
195     if (!cse.hasCandidateShape(op)) return false;
196     return isCseCandidateWithTraits(op, traits, storage);
197 }
198 
199 fn isCseCandidateWithTraits(
200     op: *ir.Operation,
201     traits: ir.OperationTraits,
202     storage: []ir.interfaces.effects.Fact,
203 ) bool {
204     if (traits.is_terminator) return false;
205     if (op.hasInterface(ir.interfaces.SymbolOpInterface)) return false;
206     if (op.hasInterface(ir.interfaces.CallOpInterface)) return false;
207     const declaration = ir.interfaces.effects.collectInto(op, storage) catch return false;
208     return ir.interfaces.effects.repeatableExpression(declaration);
209 }
210 
211 fn isCommutativeForCse(op: *ir.Operation, traits: ir.OperationTraits) bool {
212     if (traits.is_commutative) return true;
213     if (op.interface(ir.interfaces.CseOpInterface)) |iface| {
214         return iface.call(.commuteOperandsInKey, .{});
215     }
216     return false;
217 }
218 
219 fn cseEquivalenceOptions(commutative: bool) ir.OperationEquivalence.Options {
220     return .{
221         .attribute_filter = includeAttributeInCseKey,
222         .commute_operands = commutative,
223     };
224 }
225 
226 fn includeAttributeInCseKey(_: ?*const anyopaque, op: *ir.Operation, attr: ir.NamedAttribute) bool {
227     if (op.interface(ir.interfaces.CseOpInterface)) |iface| {
228         return iface.call(.includeAttributeInKey, .{ attr.name, attr.value });
229     }
230     return true;
231 }