lib/choir/src/backends/gpu/cpu/lowering.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const abi = @import("choir_abi");
3 const choir = @import("../../../root.zig");
4 const gpu = @import("../../../dialects/gpu/root.zig");
5
6 const ir = choir.ir;
7 const dialects = choir.dialects;
8 const ArithDialect = dialects.ArithDialect;
9 const BuiltinDialect = dialects.BuiltinDialect;
10 const FuncDialect = dialects.FuncDialect;
11 const GpuDialect = gpu.GpuDialect;
12 const MemrefDialect = dialects.MemrefDialect;
13 const ScfDialect = dialects.ScfDialect;
14 const Dimension = gpu.Dimension;
15
16 pub const LowerOptions = struct {
17 entry_name: []const u8,
18 vector_width: ?u32 = null,
19 };
20
21 pub fn lowerKernelToHostLoop(
22 allocator: std.mem.Allocator,
23 kernel_module: *ir.Operation,
24 options: LowerOptions,
25 ) abi.Error!*ir.Operation {
26 const source = ir.inspection.functionDefinitionByName(kernel_module, options.entry_name) orelse return error.InvalidArtifact;
27 const source_func = FuncDialect.FuncOp{ .op = source };
28 if (source_func.getNumResults() != 0) return error.UnsupportedOperation;
29
30 const host_module = BuiltinDialect.ModuleOp.create(kernel_module.getContext(), source.location) catch |err| return loweringError(err);
31 errdefer host_module.op.erase();
32
33 const source_args = source_func.getArguments();
34 const index_type = ArithDialect.getIndexType(kernel_module.getContext()) catch |err| return loweringError(err);
35 const buffer_arg_count = try countBufferArguments(kernel_module.getContext(), source_args);
36 const host_arg_types = allocator.alloc(ir.Type, source_args.len + abi.launch_shape_argument_count) catch return error.OutOfMemory;
37 defer allocator.free(host_arg_types);
38 var buffer_arg_index: usize = 0;
39 var scalar_arg_index: usize = buffer_arg_count;
40 for (source_args) |argument| {
41 const index = if (try isBufferType(kernel_module.getContext(), argument.type)) index: {
42 const current = buffer_arg_index;
43 buffer_arg_index += 1;
44 break :index current;
45 } else index: {
46 const current = scalar_arg_index;
47 scalar_arg_index += 1;
48 break :index current;
49 };
50 host_arg_types[index] = cpuBoundaryType(kernel_module.getContext(), argument.type) catch |err| return loweringError(err);
51 }
52 for (0..abi.launch_shape_argument_count) |index| {
53 host_arg_types[source_args.len + index] = index_type;
54 }
55
56 const host_func = FuncDialect.FuncOp.create(
57 kernel_module.getContext(),
58 source.location,
59 options.entry_name,
60 host_arg_types,
61 &.{},
62 ) catch |err| return loweringError(err);
63 host_module.getBodyBlock().addOperation(host_func.op) catch |err| return loweringError(err);
64
65 const entry = host_func.getEntryBlock();
66 const host_args = host_func.getArguments();
67 const mapped_source_args = allocator.alloc(*ir.Value, source_args.len) catch return error.OutOfMemory;
68 defer allocator.free(mapped_source_args);
69 buffer_arg_index = 0;
70 scalar_arg_index = buffer_arg_count;
71 for (source_args, 0..) |source_arg, source_index| {
72 const index = if (try isBufferType(kernel_module.getContext(), source_arg.type)) index: {
73 const current = buffer_arg_index;
74 buffer_arg_index += 1;
75 break :index current;
76 } else index: {
77 const current = scalar_arg_index;
78 scalar_arg_index += 1;
79 break :index current;
80 };
81 const boundary_value = try cpuBoundaryValue(kernel_module.getContext(), source.location, entry, source_arg.type, host_args[index]);
82 mapped_source_args[source_index] = boundary_value;
83 }
84
85 const source_block = source_func.getEntryBlock();
86 const vectorization = try vectorizationForBlock(allocator, kernel_module.getContext(), source_block, options.vector_width);
87
88 const zero = createIndexConstant(kernel_module.getContext(), source.location, 0) catch |err| return loweringError(err);
89 entry.addOperation(zero.op) catch |err| return loweringError(err);
90 const one = createIndexConstant(kernel_module.getContext(), source.location, 1) catch |err| return loweringError(err);
91 entry.addOperation(one.op) catch |err| return loweringError(err);
92
93 const upper = host_args[source_args.len];
94 const grid_dims = [3]*ir.Value{
95 host_args[source_args.len + 1],
96 host_args[source_args.len + 2],
97 host_args[source_args.len + 3],
98 };
99 const threadgroup_dims = [3]*ir.Value{
100 host_args[source_args.len + 4],
101 host_args[source_args.len + 5],
102 host_args[source_args.len + 6],
103 };
104 const gpu_usage = collectGpuIndexUsage(source_block);
105
106 if (vectorization) |plan| {
107 const vector_step = createIndexConstant(kernel_module.getContext(), source.location, plan.width) catch |err| return loweringError(err);
108 entry.addOperation(vector_step.op) catch |err| return loweringError(err);
109 const remainder = ArithDialect.RemOp.create(kernel_module.getContext(), source.location, upper, vector_step.getResult()) catch |err| return loweringError(err);
110 entry.addOperation(remainder.op) catch |err| return loweringError(err);
111 const vector_upper = ArithDialect.SubOp.create(kernel_module.getContext(), source.location, upper, remainder.getResult()) catch |err| return loweringError(err);
112 entry.addOperation(vector_upper.op) catch |err| return loweringError(err);
113
114 const vector_loop = ScfDialect.ForOp.create(
115 kernel_module.getContext(),
116 source.location,
117 zero.getResult(),
118 vector_upper.getResult(),
119 vector_step.getResult(),
120 &.{},
121 &.{},
122 ) catch |err| return loweringError(err);
123 entry.addOperation(vector_loop.op) catch |err| return loweringError(err);
124 try cloneKernelBlockIntoLoop(
125 allocator,
126 kernel_module.getContext(),
127 source.location,
128 source_block,
129 source_args,
130 mapped_source_args,
131 vector_loop,
132 grid_dims,
133 threadgroup_dims,
134 gpu_usage,
135 plan,
136 );
137
138 const tail_loop = ScfDialect.ForOp.create(
139 kernel_module.getContext(),
140 source.location,
141 vector_upper.getResult(),
142 upper,
143 one.getResult(),
144 &.{},
145 &.{},
146 ) catch |err| return loweringError(err);
147 entry.addOperation(tail_loop.op) catch |err| return loweringError(err);
148 try cloneKernelBlockIntoLoop(
149 allocator,
150 kernel_module.getContext(),
151 source.location,
152 source_block,
153 source_args,
154 mapped_source_args,
155 tail_loop,
156 grid_dims,
157 threadgroup_dims,
158 gpu_usage,
159 null,
160 );
161 } else {
162 const loop = ScfDialect.ForOp.create(
163 kernel_module.getContext(),
164 source.location,
165 zero.getResult(),
166 upper,
167 one.getResult(),
168 &.{},
169 &.{},
170 ) catch |err| return loweringError(err);
171 entry.addOperation(loop.op) catch |err| return loweringError(err);
172 try cloneKernelBlockIntoLoop(
173 allocator,
174 kernel_module.getContext(),
175 source.location,
176 source_block,
177 source_args,
178 mapped_source_args,
179 loop,
180 grid_dims,
181 threadgroup_dims,
182 gpu_usage,
183 null,
184 );
185 }
186
187 const ret = FuncDialect.ReturnOp.create(kernel_module.getContext(), source.location, &.{}) catch |err| return loweringError(err);
188 entry.addOperation(ret.op) catch |err| return loweringError(err);
189
190 ir.verifyOperation(host_module.op, ir.verify.default_options) catch return error.CompilationFailed;
191 return host_module.op;
192 }
193
194 fn cloneKernelBlockIntoLoop(
195 allocator: std.mem.Allocator,
196 ctx: *ir.Context,
197 loc: ir.Location,
198 source_block: *ir.Block,
199 source_args: []*ir.Value,
200 mapped_source_args: []*ir.Value,
201 loop: ScfDialect.ForOp,
202 grid_dims: [3]*ir.Value,
203 threadgroup_dims: [3]*ir.Value,
204 gpu_usage: GpuIndexUsage,
205 vectorization: ?VectorizationPlan,
206 ) abi.Error!void {
207 var mapping = ir.Mapping.init(allocator);
208 defer mapping.deinit();
209 try seedSourceArgumentMapping(&mapping, source_args, mapped_source_args);
210
211 const loop_body = loop.getBodyBlock();
212 mapping.mapBlock(source_block, loop_body) catch return error.OutOfMemory;
213 const gpu_values = try createGpuIndexValues(
214 ctx,
215 loc,
216 loop_body,
217 loop.getInductionVar(),
218 grid_dims,
219 threadgroup_dims,
220 gpu_usage,
221 );
222 if (vectorization) |plan| {
223 try cloneBlockIntoVectorHostLoop(
224 source_block,
225 loop_body,
226 &mapping,
227 gpu_values,
228 plan,
229 true,
230 );
231 } else {
232 try cloneBlockIntoHostLoop(
233 source_block,
234 loop_body,
235 &mapping,
236 gpu_values,
237 true,
238 );
239 }
240 }
241
242 fn seedSourceArgumentMapping(
243 mapping: *ir.Mapping,
244 source_args: []*ir.Value,
245 mapped_source_args: []*ir.Value,
246 ) abi.Error!void {
247 if (source_args.len != mapped_source_args.len) return error.InvalidArtifact;
248 for (source_args, mapped_source_args) |source_arg, mapped_source_arg| {
249 mapping.mapValue(source_arg, mapped_source_arg) catch return error.OutOfMemory;
250 }
251 }
252
253 fn createIndexConstant(ctx: *ir.Context, loc: ir.Location, value: i64) !ArithDialect.ConstantOp {
254 const index_type = try ArithDialect.getIndexType(ctx);
255 return try ArithDialect.ConstantOp.createInt(ctx, loc, index_type, value);
256 }
257
258 fn cpuBoundaryType(ctx: *ir.Context, typ: ir.Type) !ir.Type {
259 if (try ctx.getTypeParamPayload(typ, MemrefDialect.MemrefTypePayload)) |payload| {
260 const element_type = payload.element_type orelse return error.InvalidMemrefElementType;
261 const attrs = MemrefDialect.MemrefTypeAttrs{
262 .alignment = payload.alignment,
263 .exclusive = payload.exclusive,
264 .indexing = payload.indexing,
265 };
266 if (payload.size) |size| {
267 return try MemrefDialect.getMemrefType1DWithAttrs(ctx, size, element_type, .host, attrs);
268 }
269 return try MemrefDialect.getMemrefTypeDynamicWithAttrs(ctx, element_type, .host, attrs);
270 }
271 return switch (cpuScalarBoundaryKind(typ) orelse return typ) {
272 .f32 => try ArithDialect.getScalarType(ctx, .i32),
273 .f64 => try ArithDialect.getScalarType(ctx, .i64),
274 else => typ,
275 };
276 }
277
278 fn cpuBoundaryValue(
279 ctx: *ir.Context,
280 loc: ir.Location,
281 entry: *ir.Block,
282 source_type: ir.Type,
283 boundary_value: *ir.Value,
284 ) abi.Error!*ir.Value {
285 switch (cpuScalarBoundaryKind(source_type) orelse return boundary_value) {
286 .f32, .f64 => {
287 const cast = ArithDialect.BitcastOp.create(ctx, loc, boundary_value, source_type) catch |err| return loweringError(err);
288 entry.addOperation(cast.op) catch |err| return loweringError(err);
289 return cast.getResult();
290 },
291 else => return boundary_value,
292 }
293 }
294
295 fn cpuScalarBoundaryKind(typ: ir.Type) ?dialects.arith.ScalarKind {
296 const type_name = typ.getDialectTypeName() orelse return null;
297 return dialects.arith.scalarKindFromTypeName(type_name);
298 }
299
300 fn countBufferArguments(ctx: *ir.Context, args: []*ir.Value) abi.Error!usize {
301 var count: usize = 0;
302 for (args) |argument| {
303 if (try isBufferType(ctx, argument.type)) count += 1;
304 }
305 return count;
306 }
307
308 fn isBufferType(ctx: *ir.Context, typ: ir.Type) abi.Error!bool {
309 return (ctx.getTypeParamPayload(typ, MemrefDialect.MemrefTypePayload) catch |err| return loweringError(err)) != null;
310 }
311
312 const VectorizationPlan = struct {
313 width: u32,
314 };
315
316 fn vectorizationForBlock(
317 allocator: std.mem.Allocator,
318 ctx: *ir.Context,
319 block: *ir.Block,
320 width: ?u32,
321 ) abi.Error!?VectorizationPlan {
322 const requested = width orelse return null;
323 if (requested < 2) return null;
324 if (!try canVectorizeBlock(allocator, ctx, block, requested)) return null;
325 return .{ .width = requested };
326 }
327
328 fn canVectorizeBlock(
329 allocator: std.mem.Allocator,
330 ctx: *ir.Context,
331 block: *ir.Block,
332 width: u32,
333 ) abi.Error!bool {
334 var index_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};
335 defer index_values.deinit(allocator);
336 var vector_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};
337 defer vector_values.deinit(allocator);
338 var scalar_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};
339 defer scalar_values.deinit(allocator);
340 var saw_load = false;
341 var saw_store = false;
342
343 for (block.arguments.items) |argument| {
344 if ((try packedCpuVectorTypeForScalarType(ctx, argument.type, width)) != null) {
345 scalar_values.put(allocator, argument, {}) catch return error.OutOfMemory;
346 }
347 }
348
349 var iter = block.getOperations();
350 while (iter.next()) |op| {
351 if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) continue;
352
353 if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {
354 if ((gpuDimension(op) orelse return false) != .x) return false;
355 const result = op.getResult(0) orelse return false;
356 index_values.put(allocator, result, {}) catch return error.OutOfMemory;
357 continue;
358 }
359
360 if (std.mem.eql(u8, op.name.getDialectNamespace(), GpuDialect.name)) return false;
361 if (op.successors.items.len != 0 or op.regions.items.len != 0) return false;
362
363 if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name)) {
364 const load = MemrefDialect.LoadOp{ .op = op };
365 if (!index_values.contains(load.getIndex())) return false;
366 _ = try packedCpuVectorTypeForScalarType(ctx, load.getResult().type, width) orelse return false;
367 vector_values.put(allocator, load.getResult(), {}) catch return error.OutOfMemory;
368 saw_load = true;
369 continue;
370 }
371
372 if (std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name)) {
373 const store = MemrefDialect.StoreOp{ .op = op };
374 if (!index_values.contains(store.getIndex())) return false;
375 if (!vector_values.contains(store.getValue())) return false;
376 saw_store = true;
377 continue;
378 }
379
380 if (std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) {
381 const result = op.getResult(0) orelse return false;
382 if ((try packedCpuVectorTypeForScalarType(ctx, result.type, width)) == null) return false;
383 scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;
384 continue;
385 }
386
387 if (isVectorizableBinaryOp(op)) {
388 const lhs = op.getOperand(0) orelse return false;
389 const rhs = op.getOperand(1) orelse return false;
390 const result = op.getResult(0) orelse return false;
391 _ = try packedCpuVectorTypeForScalarType(ctx, result.type, width) orelse return false;
392 const lhs_vector = vector_values.contains(lhs);
393 const rhs_vector = vector_values.contains(rhs);
394 if (!lhs_vector and !scalar_values.contains(lhs)) return false;
395 if (!rhs_vector and !scalar_values.contains(rhs)) return false;
396 if (lhs_vector or rhs_vector) {
397 vector_values.put(allocator, result, {}) catch return error.OutOfMemory;
398 } else {
399 scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;
400 }
401 continue;
402 }
403
404 if (isVectorizableUnaryOp(op)) {
405 const input = op.getOperand(0) orelse return false;
406 const result = op.getResult(0) orelse return false;
407 _ = try packedCpuVectorTypeForScalarType(ctx, result.type, width) orelse return false;
408 if (vector_values.contains(input)) {
409 vector_values.put(allocator, result, {}) catch return error.OutOfMemory;
410 } else if (scalar_values.contains(input)) {
411 scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;
412 } else {
413 return false;
414 }
415 continue;
416 }
417
418 return false;
419 }
420
421 return saw_load and saw_store;
422 }
423
424 fn cloneBlockIntoHostLoop(
425 source: *ir.Block,
426 dest: *ir.Block,
427 mapping: *ir.Mapping,
428 gpu_values: GpuIndexValues,
429 function_body: bool,
430 ) abi.Error!void {
431 var iter = source.getOperations();
432 while (iter.next()) |op| {
433 if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) {
434 if (!function_body) return error.UnsupportedOperation;
435 const yield_op = ScfDialect.YieldOp.create(op.getContext(), op.location, &.{}) catch |err| return loweringError(err);
436 dest.addOperation(yield_op.op) catch |err| return loweringError(err);
437 continue;
438 }
439
440 if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {
441 const global_index = gpu_values.global.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
442 const result = op.getResult(0) orelse return error.UnsupportedOperation;
443 mapping.mapValue(result, global_index) catch return error.OutOfMemory;
444 continue;
445 }
446
447 if (std.mem.eql(u8, op.name.name, GpuDialect.ThreadIdxOp.operation_name)) {
448 const thread_index = gpu_values.thread.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
449 const result = op.getResult(0) orelse return error.UnsupportedOperation;
450 mapping.mapValue(result, thread_index) catch return error.OutOfMemory;
451 continue;
452 }
453
454 if (std.mem.eql(u8, op.name.name, GpuDialect.BlockIdxOp.operation_name)) {
455 const block_index = gpu_values.block.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
456 const result = op.getResult(0) orelse return error.UnsupportedOperation;
457 mapping.mapValue(result, block_index) catch return error.OutOfMemory;
458 continue;
459 }
460
461 if (std.mem.eql(u8, op.name.name, GpuDialect.BlockDimOp.operation_name)) {
462 const block_dim = gpu_values.threadgroupValue(gpuDimension(op) orelse return error.UnsupportedOperation);
463 const result = op.getResult(0) orelse return error.UnsupportedOperation;
464 mapping.mapValue(result, block_dim) catch return error.OutOfMemory;
465 continue;
466 }
467
468 if (std.mem.eql(u8, op.name.name, GpuDialect.GridDimOp.operation_name)) {
469 const grid_dim = gpu_values.gridValue(gpuDimension(op) orelse return error.UnsupportedOperation);
470 const result = op.getResult(0) orelse return error.UnsupportedOperation;
471 mapping.mapValue(result, grid_dim) catch return error.OutOfMemory;
472 continue;
473 }
474
475 if (std.mem.eql(u8, op.name.getDialectNamespace(), GpuDialect.name)) {
476 return error.UnsupportedOperation;
477 }
478
479 if (op.successors.items.len != 0) return error.UnsupportedOperation;
480
481 const has_regions = op.regions.items.len != 0;
482 const cloned = op.cloneWithoutRegionsMapped(mapping, .{ .clone_operands = !has_regions }) catch |err| return loweringError(err);
483 errdefer cloned.erase();
484 dest.addOperation(cloned) catch |err| return loweringError(err);
485 for (op.regions.items, 0..) |*source_region, region_index| {
486 const dest_region = &cloned.regions.items[region_index];
487 var source_block = source_region.blocks.head;
488 while (source_block) |block| : (source_block = block.next) {
489 const cloned_block = dest_region.addBlock() catch |err| return loweringError(err);
490 mapping.mapBlock(block, cloned_block) catch return error.OutOfMemory;
491 cloned_block.arguments.ensureTotalCapacity(cloned_block.allocator, block.arguments.items.len) catch return error.OutOfMemory;
492 for (block.arguments.items) |argument| {
493 const cloned_argument = cloned_block.addArgument(argument.type, .unknown) catch |err| return loweringError(err);
494 mapping.mapValue(argument, cloned_argument) catch return error.OutOfMemory;
495 }
496 }
497
498 source_block = source_region.blocks.head;
499 while (source_block) |block| : (source_block = block.next) {
500 const cloned_block = mapping.lookupBlock(block) orelse return error.UnsupportedOperation;
501 try cloneBlockIntoHostLoop(block, cloned_block, mapping, gpu_values, false);
502 }
503 }
504 if (has_regions) try replaceOperandsMapped(op, cloned, mapping);
505 }
506 }
507
508 fn cloneBlockIntoVectorHostLoop(
509 source: *ir.Block,
510 dest: *ir.Block,
511 mapping: *ir.Mapping,
512 gpu_values: GpuIndexValues,
513 plan: VectorizationPlan,
514 function_body: bool,
515 ) abi.Error!void {
516 var iter = source.getOperations();
517 while (iter.next()) |op| {
518 if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) {
519 if (!function_body) return error.UnsupportedOperation;
520 const yield_op = ScfDialect.YieldOp.create(op.getContext(), op.location, &.{}) catch |err| return loweringError(err);
521 dest.addOperation(yield_op.op) catch |err| return loweringError(err);
522 continue;
523 }
524
525 if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {
526 if ((gpuDimension(op) orelse return error.UnsupportedOperation) != .x) return error.UnsupportedOperation;
527 const global_index = gpu_values.global.x orelse return error.UnsupportedOperation;
528 const result = op.getResult(0) orelse return error.UnsupportedOperation;
529 mapping.mapValue(result, global_index) catch return error.OutOfMemory;
530 continue;
531 }
532
533 if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name)) {
534 try cloneVectorLoad(op, dest, mapping, plan.width);
535 continue;
536 }
537
538 if (std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name)) {
539 try cloneVectorStore(op, dest, mapping);
540 continue;
541 }
542
543 if (std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) {
544 try cloneScalarOperationIntoVectorLoop(op, dest, mapping);
545 continue;
546 }
547
548 if (isVectorizableBinaryOp(op)) {
549 try cloneVectorBinaryOp(op, dest, mapping, plan.width);
550 continue;
551 }
552
553 if (isVectorizableUnaryOp(op)) {
554 try cloneVectorUnaryOp(op, dest, mapping, plan.width);
555 continue;
556 }
557
558 return error.UnsupportedOperation;
559 }
560 }
561
562 fn cloneVectorLoad(
563 op: *ir.Operation,
564 dest: *ir.Block,
565 mapping: *ir.Mapping,
566 width: u32,
567 ) abi.Error!void {
568 const load = MemrefDialect.LoadOp{ .op = op };
569 const memref = mapping.lookupValue(load.getMemref()) orelse return error.UnsupportedOperation;
570 const index = mapping.lookupValue(load.getIndex()) orelse return error.UnsupportedOperation;
571 const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), load.getResult().type, width) orelse return error.UnsupportedOperation;
572 var cloned = MemrefDialect.LoadOp.create(op.getContext(), op.location, memref, index, result_type) catch |err| return loweringError(err);
573 dest.addOperation(cloned.op) catch |err| return loweringError(err);
574 mapping.mapValue(load.getResult(), cloned.getResult()) catch return error.OutOfMemory;
575 }
576
577 fn cloneVectorStore(
578 op: *ir.Operation,
579 dest: *ir.Block,
580 mapping: *ir.Mapping,
581 ) abi.Error!void {
582 const store = MemrefDialect.StoreOp{ .op = op };
583 const value = mapping.lookupValue(store.getValue()) orelse return error.UnsupportedOperation;
584 const memref = mapping.lookupValue(store.getMemref()) orelse return error.UnsupportedOperation;
585 const index = mapping.lookupValue(store.getIndex()) orelse return error.UnsupportedOperation;
586 const cloned = MemrefDialect.StoreOp.create(op.getContext(), op.location, value, memref, index) catch |err| return loweringError(err);
587 dest.addOperation(cloned.op) catch |err| return loweringError(err);
588 }
589
590 fn cloneVectorBinaryOp(
591 op: *ir.Operation,
592 dest: *ir.Block,
593 mapping: *ir.Mapping,
594 width: u32,
595 ) abi.Error!void {
596 const lhs = mapping.lookupValue(op.getOperand(0) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
597 const rhs = mapping.lookupValue(op.getOperand(1) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
598 const source_result = op.getResult(0) orelse return error.UnsupportedOperation;
599 const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), source_result.type, width) orelse return error.UnsupportedOperation;
600 if (!lhs.type.eql(result_type) and !rhs.type.eql(result_type)) {
601 try cloneScalarOperationIntoVectorLoop(op, dest, mapping);
602 return;
603 }
604
605 const vector_lhs = try vectorOperandForMappedValue(op, dest, lhs, result_type, width);
606 const vector_rhs = try vectorOperandForMappedValue(op, dest, rhs, result_type, width);
607
608 var state = ir.Operation.State.init(op.name.name, op.location);
609 state.addOperands(&.{ vector_lhs, vector_rhs });
610 state.addTypes(&.{result_type});
611 const cloned = op.getContext().createOperation(state) catch |err| return loweringError(err);
612 errdefer cloned.erase();
613 dest.addOperation(cloned) catch |err| return loweringError(err);
614 const result = cloned.getResult(0) orelse return error.UnsupportedOperation;
615 mapping.mapValue(source_result, result) catch return error.OutOfMemory;
616 }
617
618 fn cloneVectorUnaryOp(
619 op: *ir.Operation,
620 dest: *ir.Block,
621 mapping: *ir.Mapping,
622 width: u32,
623 ) abi.Error!void {
624 const input = mapping.lookupValue(op.getOperand(0) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;
625 const source_result = op.getResult(0) orelse return error.UnsupportedOperation;
626 const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), source_result.type, width) orelse return error.UnsupportedOperation;
627 if (!input.type.eql(result_type)) {
628 try cloneScalarOperationIntoVectorLoop(op, dest, mapping);
629 return;
630 }
631
632 var state = ir.Operation.State.init(op.name.name, op.location);
633 state.addOperands(&.{input});
634 state.addTypes(&.{result_type});
635 const cloned = op.getContext().createOperation(state) catch |err| return loweringError(err);
636 errdefer cloned.erase();
637 dest.addOperation(cloned) catch |err| return loweringError(err);
638 const result = cloned.getResult(0) orelse return error.UnsupportedOperation;
639 mapping.mapValue(source_result, result) catch return error.OutOfMemory;
640 }
641
642 fn cloneScalarOperationIntoVectorLoop(
643 op: *ir.Operation,
644 dest: *ir.Block,
645 mapping: *ir.Mapping,
646 ) abi.Error!void {
647 const cloned = op.cloneWithoutRegionsMapped(mapping, .{ .clone_operands = true }) catch |err| return loweringError(err);
648 errdefer cloned.erase();
649 dest.addOperation(cloned) catch |err| return loweringError(err);
650 }
651
652 fn vectorOperandForMappedValue(
653 op: *ir.Operation,
654 dest: *ir.Block,
655 mapped: *ir.Value,
656 result_type: ir.Type,
657 width: u32,
658 ) abi.Error!*ir.Value {
659 if (mapped.type.eql(result_type)) return mapped;
660 const vector_type = try packedCpuVectorTypeForScalarType(op.getContext(), mapped.type, width) orelse return error.UnsupportedOperation;
661 if (!vector_type.eql(result_type)) return error.UnsupportedOperation;
662 const splat = ArithDialect.SplatOp.create(op.getContext(), op.location, mapped, result_type) catch |err| return loweringError(err);
663 dest.addOperation(splat.op) catch |err| return loweringError(err);
664 return splat.getResult();
665 }
666
667 fn isVectorizableBinaryOp(op: *const ir.Operation) bool {
668 return std.mem.eql(u8, op.name.name, ArithDialect.AddOp.operation_name) or
669 std.mem.eql(u8, op.name.name, ArithDialect.SubOp.operation_name) or
670 std.mem.eql(u8, op.name.name, ArithDialect.MulOp.operation_name) or
671 std.mem.eql(u8, op.name.name, ArithDialect.DivOp.operation_name) or
672 std.mem.eql(u8, op.name.name, ArithDialect.MinOp.operation_name) or
673 std.mem.eql(u8, op.name.name, ArithDialect.MaxOp.operation_name) or
674 std.mem.eql(u8, op.name.name, ArithDialect.AndOp.operation_name) or
675 std.mem.eql(u8, op.name.name, ArithDialect.OrOp.operation_name) or
676 std.mem.eql(u8, op.name.name, ArithDialect.XorOp.operation_name) or
677 std.mem.eql(u8, op.name.name, ArithDialect.ShlOp.operation_name) or
678 std.mem.eql(u8, op.name.name, ArithDialect.ShrOp.operation_name) or
679 std.mem.eql(u8, op.name.name, ArithDialect.UshrOp.operation_name);
680 }
681
682 fn isVectorizableUnaryOp(op: *const ir.Operation) bool {
683 return std.mem.eql(u8, op.name.name, ArithDialect.NegOp.operation_name) or
684 std.mem.eql(u8, op.name.name, ArithDialect.NotOp.operation_name) or
685 std.mem.eql(u8, op.name.name, ArithDialect.PopCountOp.operation_name);
686 }
687
688 fn packedCpuVectorTypeForScalarType(ctx: *ir.Context, scalar_type: ir.Type, width: u32) abi.Error!?ir.Type {
689 const scalar_name = scalar_type.getDialectTypeName() orelse return null;
690 const scalar_kind = dialects.arith.scalarKindFromTypeName(scalar_name) orelse return null;
691 switch (scalar_kind) {
692 .f32, .i32, .u32 => if (width != 4) return null,
693 .f64, .i64, .u64 => if (width != 2) return null,
694 else => return null,
695 }
696 const vector_name = dialects.arith.vectorTypeNameForElement(width, scalar_name) orelse return null;
697 return ctx.getDialectTypeFromName(vector_name) catch |err| return loweringError(err);
698 }
699
700 fn replaceOperandsMapped(source: *ir.Operation, dest: *ir.Operation, mapping: *ir.Mapping) abi.Error!void {
701 var operands: []*ir.Value = &.{};
702 defer if (operands.len != 0) dest.allocator.free(operands);
703 if (source.operand_values.len != 0) {
704 operands = dest.allocator.alloc(*ir.Value, source.operand_values.len) catch return error.OutOfMemory;
705 for (source.operand_values, 0..) |operand, index| {
706 operands[index] = mapping.lookupOrDefaultValue(operand);
707 }
708 }
709 dest.replaceOperands(operands) catch |err| return loweringError(err);
710 }
711
712 const DimensionUsage = struct {
713 x: bool = false,
714 y: bool = false,
715 z: bool = false,
716
717 fn add(self: *DimensionUsage, dim: Dimension) void {
718 switch (dim) {
719 .x => self.x = true,
720 .y => self.y = true,
721 .z => self.z = true,
722 }
723 }
724
725 fn addAll(self: *DimensionUsage, other: DimensionUsage) void {
726 self.x = self.x or other.x;
727 self.y = self.y or other.y;
728 self.z = self.z or other.z;
729 }
730
731 fn contains(self: DimensionUsage, dim: Dimension) bool {
732 return switch (dim) {
733 .x => self.x,
734 .y => self.y,
735 .z => self.z,
736 };
737 }
738 };
739
740 const GpuIndexUsage = struct {
741 global: DimensionUsage = .{},
742 thread: DimensionUsage = .{},
743 block: DimensionUsage = .{},
744 block_dim: DimensionUsage = .{},
745 grid_dim: DimensionUsage = .{},
746
747 fn merge(self: *GpuIndexUsage, other: GpuIndexUsage) void {
748 self.global.addAll(other.global);
749 self.thread.addAll(other.thread);
750 self.block.addAll(other.block);
751 self.block_dim.addAll(other.block_dim);
752 self.grid_dim.addAll(other.grid_dim);
753 }
754
755 fn derived(self: GpuIndexUsage) DimensionUsage {
756 var result = DimensionUsage{};
757 result.addAll(self.global);
758 result.addAll(self.thread);
759 result.addAll(self.block);
760 return result;
761 }
762 };
763
764 const DimensionValues = struct {
765 x: ?*ir.Value = null,
766 y: ?*ir.Value = null,
767 z: ?*ir.Value = null,
768
769 fn value(self: DimensionValues, dim: Dimension) ?*ir.Value {
770 return switch (dim) {
771 .x => self.x,
772 .y => self.y,
773 .z => self.z,
774 };
775 }
776
777 fn set(self: *DimensionValues, dim: Dimension, new_value: *ir.Value) void {
778 switch (dim) {
779 .x => self.x = new_value,
780 .y => self.y = new_value,
781 .z => self.z = new_value,
782 }
783 }
784 };
785
786 const GpuIndexValues = struct {
787 grid: [3]*ir.Value,
788 threadgroup: [3]*ir.Value,
789 global: DimensionValues = .{},
790 thread: DimensionValues = .{},
791 block: DimensionValues = .{},
792
793 fn gridValue(self: GpuIndexValues, dim: Dimension) *ir.Value {
794 return self.grid[dimensionIndex(dim)];
795 }
796
797 fn threadgroupValue(self: GpuIndexValues, dim: Dimension) *ir.Value {
798 return self.threadgroup[dimensionIndex(dim)];
799 }
800 };
801
802 fn collectGpuIndexUsage(block: *ir.Block) GpuIndexUsage {
803 var usage = GpuIndexUsage{};
804 var iter = block.getOperations();
805 while (iter.next()) |op| {
806 if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {
807 if (gpuDimension(op)) |dim| usage.global.add(dim);
808 } else if (std.mem.eql(u8, op.name.name, GpuDialect.ThreadIdxOp.operation_name)) {
809 if (gpuDimension(op)) |dim| usage.thread.add(dim);
810 } else if (std.mem.eql(u8, op.name.name, GpuDialect.BlockIdxOp.operation_name)) {
811 if (gpuDimension(op)) |dim| usage.block.add(dim);
812 } else if (std.mem.eql(u8, op.name.name, GpuDialect.BlockDimOp.operation_name)) {
813 if (gpuDimension(op)) |dim| usage.block_dim.add(dim);
814 } else if (std.mem.eql(u8, op.name.name, GpuDialect.GridDimOp.operation_name)) {
815 if (gpuDimension(op)) |dim| usage.grid_dim.add(dim);
816 }
817 for (op.regions.items) |*region| {
818 var block_iter = region.getBlocks();
819 while (block_iter.next()) |child_block| {
820 usage.merge(collectGpuIndexUsage(child_block));
821 }
822 }
823 }
824 return usage;
825 }
826
827 fn createGpuIndexValues(
828 ctx: *ir.Context,
829 loc: ir.Location,
830 dest: *ir.Block,
831 linear: *ir.Value,
832 grid: [3]*ir.Value,
833 threadgroup: [3]*ir.Value,
834 usage: GpuIndexUsage,
835 ) abi.Error!GpuIndexValues {
836 var values = GpuIndexValues{
837 .grid = grid,
838 .threadgroup = threadgroup,
839 };
840 const derived = usage.derived();
841 var y_linear: ?*ir.Value = null;
842 var x_extent: ?*ir.Value = null;
843 var y_extent: ?*ir.Value = null;
844
845 if (derived.x or derived.y or derived.z) {
846 x_extent = try createProduct(ctx, loc, dest, grid[0], threadgroup[0]);
847 }
848 if (derived.y or derived.z) {
849 y_extent = try createProduct(ctx, loc, dest, grid[1], threadgroup[1]);
850 }
851
852 if (derived.x) {
853 const x = ArithDialect.RemOp.create(ctx, loc, linear, x_extent.?) catch |err| return loweringError(err);
854 dest.addOperation(x.op) catch |err| return loweringError(err);
855 values.global.x = x.getResult();
856 try createThreadBlockValues(ctx, loc, dest, &values, .x, x.getResult());
857 }
858
859 if (derived.y or derived.z) {
860 const div = ArithDialect.DivOp.create(ctx, loc, linear, x_extent.?) catch |err| return loweringError(err);
861 dest.addOperation(div.op) catch |err| return loweringError(err);
862 y_linear = div.getResult();
863 }
864
865 if (derived.y) {
866 const y = ArithDialect.RemOp.create(ctx, loc, y_linear.?, y_extent.?) catch |err| return loweringError(err);
867 dest.addOperation(y.op) catch |err| return loweringError(err);
868 values.global.y = y.getResult();
869 try createThreadBlockValues(ctx, loc, dest, &values, .y, y.getResult());
870 }
871
872 if (derived.z) {
873 const xy_extent = try createProduct(ctx, loc, dest, x_extent.?, y_extent.?);
874 const z = ArithDialect.DivOp.create(ctx, loc, linear, xy_extent) catch |err| return loweringError(err);
875 dest.addOperation(z.op) catch |err| return loweringError(err);
876 values.global.z = z.getResult();
877 try createThreadBlockValues(ctx, loc, dest, &values, .z, z.getResult());
878 }
879
880 return values;
881 }
882
883 fn createThreadBlockValues(
884 ctx: *ir.Context,
885 loc: ir.Location,
886 dest: *ir.Block,
887 values: *GpuIndexValues,
888 dim: Dimension,
889 global_value: *ir.Value,
890 ) abi.Error!void {
891 const thread_extent = values.threadgroupValue(dim);
892 const thread = ArithDialect.RemOp.create(ctx, loc, global_value, thread_extent) catch |err| return loweringError(err);
893 dest.addOperation(thread.op) catch |err| return loweringError(err);
894 values.thread.set(dim, thread.getResult());
895
896 const block = ArithDialect.DivOp.create(ctx, loc, global_value, thread_extent) catch |err| return loweringError(err);
897 dest.addOperation(block.op) catch |err| return loweringError(err);
898 values.block.set(dim, block.getResult());
899 }
900
901 fn createProduct(
902 ctx: *ir.Context,
903 loc: ir.Location,
904 dest: *ir.Block,
905 lhs: *ir.Value,
906 rhs: *ir.Value,
907 ) abi.Error!*ir.Value {
908 const product = ArithDialect.MulOp.create(ctx, loc, lhs, rhs) catch |err| return loweringError(err);
909 dest.addOperation(product.op) catch |err| return loweringError(err);
910 return product.getResult();
911 }
912
913 fn dimensionIndex(dim: Dimension) usize {
914 return switch (dim) {
915 .x => 0,
916 .y => 1,
917 .z => 2,
918 };
919 }
920
921 fn gpuDimension(op: *const ir.Operation) ?Dimension {
922 const attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;
923 return Dimension.fromString(attr.payload);
924 }
925
926 fn loweringError(err: anyerror) abi.Error {
927 return switch (err) {
928 error.OutOfMemory => error.OutOfMemory,
929 else => error.CompilationFailed,
930 };
931 }