lib/accy/src/kernel/model/logical/builder.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const choir_abi = @import("choir_abi");
2 const core = @import("../core/root.zig");
3
4 const base = core.builder;
5 const domain = @import("domain/root.zig");
6 const kernel_domain = core.domain;
7 const typed = core.typed;
8 const vector = core.vector;
9 const view = core.view;
10
11 const DType = choir_abi.DType;
12
13 pub fn Builder(comptime Inner: type, comptime Schedule: type) type {
14 return struct {
15 token: *anyopaque,
16 schedule: Schedule,
17
18 const Self = @This();
19
20 pub fn init(source: *Inner, schedule: Schedule) Self {
21 return .{
22 .token = source,
23 .schedule = schedule,
24 };
25 }
26
27 pub fn index1D(self: *Self, name: []const u8, extent: u64) !kernel_domain.Index1D {
28 return self.schedule.index1D(self.owner(), domain.axis(name, extent));
29 }
30
31 pub fn index2D(self: *Self, shape: domain.Domain2D) !kernel_domain.Index2D {
32 return self.schedule.index2D(self.owner(), shape);
33 }
34
35 pub fn index3D(self: *Self, shape: domain.Domain3D) !kernel_domain.Index3D {
36 return self.schedule.index3D(self.owner(), shape);
37 }
38
39 pub fn guardDo(self: *Self, condition: anytype, context: anytype, comptime body: anytype) !void {
40 var active_guard = try self.owner().guard(condition);
41 errdefer active_guard.abort();
42 try body(self, context);
43 try active_guard.leave();
44 }
45
46 pub fn guardIndexDo(self: *Self, index: kernel_domain.Index1D, context: anytype, comptime body: anytype) !void {
47 var active_guard = try self.owner().guardIndex(index);
48 errdefer active_guard.abort();
49 try body(self, index, context);
50 try active_guard.leave();
51 }
52
53 pub fn guardIndex2DDo(self: *Self, index: kernel_domain.Index2D, context: anytype, comptime body: anytype) !void {
54 var x_guard = try self.owner().guardIndex(index.x);
55 errdefer x_guard.abort();
56 var y_guard = try self.owner().guardIndex(index.y);
57 errdefer y_guard.abort();
58 try body(self, index, context);
59 try y_guard.leave();
60 try x_guard.leave();
61 }
62
63 pub fn guardIndex3DDo(self: *Self, index: kernel_domain.Index3D, context: anytype, comptime body: anytype) !void {
64 var x_guard = try self.owner().guardIndex(index.x);
65 errdefer x_guard.abort();
66 var y_guard = try self.owner().guardIndex(index.y);
67 errdefer y_guard.abort();
68 var z_guard = try self.owner().guardIndex(index.z);
69 errdefer z_guard.abort();
70 try body(self, index, context);
71 try z_guard.leave();
72 try y_guard.leave();
73 try x_guard.leave();
74 }
75
76 pub fn forEach1D(self: *Self, name: []const u8, extent: u64, context: anytype, comptime body: anytype) !kernel_domain.Index1D {
77 const index = try self.index1D(name, extent);
78 try self.guardIndexDo(index, context, body);
79 return index;
80 }
81
82 pub fn forEach2D(self: *Self, shape: domain.Domain2D, context: anytype, comptime body: anytype) !kernel_domain.Index2D {
83 const index = try self.index2D(shape);
84 try self.guardIndex2DDo(index, context, body);
85 return index;
86 }
87
88 pub fn forEach3D(self: *Self, shape: domain.Domain3D, context: anytype, comptime body: anytype) !kernel_domain.Index3D {
89 const index = try self.index3D(shape);
90 try self.guardIndex3DDo(index, context, body);
91 return index;
92 }
93
94 pub fn fold(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, initial: anytype, context: anytype, comptime body: anytype) !@TypeOf(initial) {
95 return self.owner().fold(lower, upper, step, initial, .{
96 .logical = self,
97 .payload = context,
98 }, struct {
99 fn each(_: anytype, iteration: anytype, acc: anytype, fold_ctx: anytype) !@TypeOf(acc) {
100 return body(fold_ctx.logical, iteration, acc, fold_ctx.payload);
101 }
102 }.each);
103 }
104
105 pub fn foldRange(self: *Self, lower: i64, upper: i64, step: i64, initial: anytype, context: anytype, comptime body: anytype) !@TypeOf(initial) {
106 const lower_value = try self.constantIndex(lower);
107 const upper_value = try self.constantIndex(upper);
108 const step_value = try self.constantIndex(step);
109 return self.fold(lower_value, upper_value, step_value, initial, context, body);
110 }
111
112 pub fn whileLoop(self: *Self, initial: anytype, context: anytype, comptime condition: anytype, comptime body: anytype) !@TypeOf(initial) {
113 return self.owner().whileLoop(initial, .{
114 .logical = self,
115 .payload = context,
116 }, struct {
117 fn keepGoing(_: anytype, carry: anytype, loop_ctx: anytype) !base.Value {
118 return typed.raw(try condition(loop_ctx.logical, carry, loop_ctx.payload));
119 }
120 }.keepGoing, struct {
121 fn each(_: anytype, carry: anytype, loop_ctx: anytype) !@TypeOf(carry) {
122 return body(loop_ctx.logical, carry, loop_ctx.payload);
123 }
124 }.each);
125 }
126
127 pub fn forScope(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, init_args: []const base.Value, result_types: []const base.Type) !ForScope(Inner) {
128 return self.owner().forScope(lower, upper, step, init_args, result_types);
129 }
130
131 pub fn whileScope(self: *Self, init_args: []const base.Value, result_types: []const base.Type) !WhileScope(Inner) {
132 return self.owner().whileScope(init_args, result_types);
133 }
134
135 pub fn forDo(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, context: anytype, comptime body: anytype) !base.For {
136 var scope = try self.forScope(lower, upper, step, &.{}, &.{});
137 errdefer scope.abort();
138 try body(self, scope.inductionVar(), context);
139 try scope.leave(&.{});
140 return scope.loop;
141 }
142
143 pub fn forValueDo(self: *Self, comptime dtype: DType, lower: base.Value, upper: base.Value, step: base.Value, context: anytype, comptime body: anytype) !base.For {
144 var scope = try self.forScope(lower, upper, step, &.{}, &.{});
145 errdefer scope.abort();
146 try body(self, try self.castValue(scope.inductionVar(), dtype), context);
147 try scope.leave(&.{});
148 return scope.loop;
149 }
150
151 pub fn forRangeDo(self: *Self, lower: i64, upper: i64, step: i64, context: anytype, comptime body: anytype) !base.For {
152 const lower_value = try self.constantIndex(lower);
153 const upper_value = try self.constantIndex(upper);
154 const step_value = try self.constantIndex(step);
155 return self.forDo(lower_value, upper_value, step_value, context, body);
156 }
157
158 pub fn forRangeValueDo(self: *Self, comptime dtype: DType, lower: i64, upper: i64, step: i64, context: anytype, comptime body: anytype) !base.For {
159 const lower_value = try self.constantIndex(lower);
160 const upper_value = try self.constantIndex(upper);
161 const step_value = try self.constantIndex(step);
162 return self.forValueDo(dtype, lower_value, upper_value, step_value, context, body);
163 }
164
165 pub fn argument(self: *Self, index: usize) base.Value {
166 return self.owner().argument(index);
167 }
168
169 pub fn globalId(self: *Self, dim: base.Dimension) !base.Value {
170 return self.owner().globalId(dim);
171 }
172
173 pub fn globalIdValue(self: *Self, comptime dtype: DType, dim: base.Dimension) !typed.Value(dtype) {
174 return self.owner().globalIdValue(dtype, dim);
175 }
176
177 pub fn laneId(self: *Self) !base.Value {
178 return self.owner().laneId();
179 }
180
181 pub fn threadId(self: *Self, dim: base.Dimension) !base.Value {
182 return self.owner().threadId(dim);
183 }
184
185 pub fn blockId(self: *Self, dim: base.Dimension) !base.Value {
186 return self.owner().blockId(dim);
187 }
188
189 pub fn blockDim(self: *Self, dim: base.Dimension) !base.Value {
190 return self.owner().blockDim(dim);
191 }
192
193 pub fn gridDim(self: *Self, dim: base.Dimension) !base.Value {
194 return self.owner().gridDim(dim);
195 }
196
197 pub fn warpReduce(self: *Self, op_kind: base.WarpOpKind, value: anytype) !@TypeOf(value) {
198 return self.owner().warpReduce(op_kind, value);
199 }
200
201 pub fn warpScan(self: *Self, op_kind: base.WarpOpKind, mode: base.WarpScanMode, value: anytype) !@TypeOf(value) {
202 return self.owner().warpScan(op_kind, mode, value);
203 }
204
205 pub fn warpId(self: *Self) !base.Value {
206 return self.owner().warpId();
207 }
208
209 pub fn mmaSync(self: *Self, shape: base.MmaShape, a: [4]base.Value, b: [2]base.Value, c: [4]base.Value) ![4]base.Value {
210 return self.owner().mmaSync(shape, a, b, c);
211 }
212
213 pub fn shuffleSync(self: *Self, mode: base.ShuffleMode, value: base.Value, lane_or_delta: base.Value) !base.Value {
214 return self.owner().shuffleSync(mode, value, lane_or_delta);
215 }
216
217 pub fn fence(self: *Self, scope: base.Scope) !void {
218 return self.owner().fence(scope);
219 }
220
221 pub fn asyncCopyShared(self: *Self, dst: base.Value, dst_index: base.Value, src: base.Value, src_index: base.Value, bytes: u32) !void {
222 return self.owner().asyncCopyShared(dst, dst_index, src, src_index, bytes);
223 }
224
225 pub fn asyncCopyCommit(self: *Self) !void {
226 return self.owner().asyncCopyCommit();
227 }
228
229 pub fn asyncCopyWait(self: *Self, groups: u32) !void {
230 return self.owner().asyncCopyWait(groups);
231 }
232
233 pub fn sharedBuffer(self: *Self, dtype: DType, size: u64) !base.Value {
234 return self.owner().sharedBuffer(dtype, size);
235 }
236
237 pub fn dynamicSharedBuffer(self: *Self, dtype: DType, size: u64, byte_offset: u64) !base.Value {
238 return self.owner().dynamicSharedBuffer(dtype, size, byte_offset);
239 }
240
241 pub fn barrier(self: *Self, scope: base.Scope) !void {
242 return self.owner().barrier(scope);
243 }
244
245 pub fn typedArgument(self: *Self, comptime dtype: DType, index: usize) typed.Value(dtype) {
246 return self.owner().typedArgument(dtype, index);
247 }
248
249 pub fn typedValue(self: *Self, comptime dtype: DType, value: base.Value) typed.Value(dtype) {
250 return self.owner().typedValue(dtype, value);
251 }
252
253 pub fn castValue(self: *Self, input: anytype, comptime dtype: DType) !typed.Value(dtype) {
254 return self.owner().castValue(input, dtype);
255 }
256
257 pub fn bufferArgument(self: *Self, comptime dtype: DType, index: usize) view.BufferView(dtype) {
258 return self.owner().bufferArgument(dtype, index);
259 }
260
261 pub fn bufferView(self: *Self, comptime dtype: DType, value: base.Value) view.BufferView(dtype) {
262 return self.owner().bufferView(dtype, value);
263 }
264
265 pub fn constantInt(self: *Self, dtype: DType, value: i64) !base.Value {
266 return self.owner().constantInt(dtype, value);
267 }
268
269 pub fn constantIndex(self: *Self, value: i64) !base.Value {
270 return self.owner().constantIndex(value);
271 }
272
273 pub fn constantFloat(self: *Self, dtype: DType, value: f64) !base.Value {
274 return self.owner().constantFloat(dtype, value);
275 }
276
277 pub fn constantBool(self: *Self, value: bool) !base.Value {
278 return self.owner().constantBool(value);
279 }
280
281 pub fn constantValue(self: *Self, comptime dtype: DType, value: anytype) !typed.Value(dtype) {
282 return self.owner().constantValue(dtype, value);
283 }
284
285 pub fn add(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
286 return self.owner().add(lhs, rhs);
287 }
288
289 pub fn sub(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
290 return self.owner().sub(lhs, rhs);
291 }
292
293 pub fn mul(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
294 return self.owner().mul(lhs, rhs);
295 }
296
297 pub fn umulhi(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
298 return self.owner().umulhi(lhs, rhs);
299 }
300
301 pub fn div(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
302 return self.owner().div(lhs, rhs);
303 }
304
305 pub fn min(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
306 return self.owner().min(lhs, rhs);
307 }
308
309 pub fn max(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
310 return self.owner().max(lhs, rhs);
311 }
312
313 pub fn and_(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
314 return self.owner().and_(lhs, rhs);
315 }
316
317 pub fn or_(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
318 return self.owner().or_(lhs, rhs);
319 }
320
321 pub fn xor(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {
322 return self.owner().xor(lhs, rhs);
323 }
324
325 pub fn not(self: *Self, input: base.Value) !base.Value {
326 return self.owner().not(input);
327 }
328
329 pub fn popcount(self: *Self, input: base.Value) !base.Value {
330 return self.owner().popcount(input);
331 }
332
333 pub fn ballotSync(self: *Self, predicate: base.Value) !base.Value {
334 return self.owner().ballotSync(predicate);
335 }
336
337 pub fn shl(self: *Self, value: base.Value, shift: base.Value) !base.Value {
338 return self.owner().shl(value, shift);
339 }
340
341 pub fn shr(self: *Self, value: base.Value, shift: base.Value) !base.Value {
342 return self.owner().shr(value, shift);
343 }
344
345 pub fn ushr(self: *Self, value: base.Value, shift: base.Value) !base.Value {
346 return self.owner().ushr(value, shift);
347 }
348
349 pub fn bitcast(self: *Self, input: base.Value, dtype: DType) !base.Value {
350 return self.owner().bitcast(input, dtype);
351 }
352
353 pub fn neg(self: *Self, input: base.Value) !base.Value {
354 return self.owner().neg(input);
355 }
356
357 pub fn abs(self: *Self, input: base.Value) !base.Value {
358 return self.owner().abs(input);
359 }
360
361 pub fn sqrt(self: *Self, input: base.Value) !base.Value {
362 return self.owner().sqrt(input);
363 }
364
365 pub fn exp(self: *Self, input: base.Value) !base.Value {
366 return self.owner().exp(input);
367 }
368
369 pub fn log(self: *Self, input: base.Value) !base.Value {
370 return self.owner().log(input);
371 }
372
373 pub fn tanh(self: *Self, input: base.Value) !base.Value {
374 return self.owner().tanh(input);
375 }
376
377 pub fn sin(self: *Self, input: base.Value) !base.Value {
378 return self.owner().sin(input);
379 }
380
381 pub fn cos(self: *Self, input: base.Value) !base.Value {
382 return self.owner().cos(input);
383 }
384
385 pub fn tan(self: *Self, input: base.Value) !base.Value {
386 return self.owner().tan(input);
387 }
388
389 pub fn floor(self: *Self, input: base.Value) !base.Value {
390 return self.owner().floor(input);
391 }
392
393 pub fn round(self: *Self, input: base.Value) !base.Value {
394 return self.owner().round(input);
395 }
396
397 pub fn trunc(self: *Self, input: base.Value) !base.Value {
398 return self.owner().trunc(input);
399 }
400
401 pub fn tf32Round(self: *Self, input: base.Value) !base.Value {
402 return self.owner().tf32Round(input);
403 }
404
405 pub fn pow(self: *Self, base_value: base.Value, exponent: base.Value) !base.Value {
406 return self.owner().pow(base_value, exponent);
407 }
408
409 pub fn atan2(self: *Self, y: base.Value, x: base.Value) !base.Value {
410 return self.owner().atan2(y, x);
411 }
412
413 pub fn fma(self: *Self, a: base.Value, b: base.Value, c: base.Value) !base.Value {
414 return self.owner().fma(a, b, c);
415 }
416
417 pub fn compare(self: *Self, predicate: base.Compare, lhs: base.Value, rhs: base.Value) !base.Value {
418 return self.owner().compare(predicate, lhs, rhs);
419 }
420
421 pub fn select(self: *Self, condition: base.Value, true_value: base.Value, false_value: base.Value) !base.Value {
422 return self.owner().select(condition, true_value, false_value);
423 }
424
425 pub fn cast(self: *Self, input: base.Value, dtype: DType) !base.Value {
426 return self.owner().cast(input, dtype);
427 }
428
429 pub fn castIndex(self: *Self, input: base.Value) !base.Value {
430 return self.owner().castIndex(input);
431 }
432
433 pub fn linearIndex(self: *Self, index: anytype) !base.Value {
434 return self.owner().linearIndex(index);
435 }
436
437 pub fn linearIndexValue(self: *Self, index: anytype, comptime dtype: DType) !typed.Value(dtype) {
438 return self.owner().linearIndexValue(index, dtype);
439 }
440
441 pub fn loadIndex(self: *Self, memref: base.Value, index: anytype) !base.Value {
442 return self.owner().loadIndex(memref, index);
443 }
444
445 pub fn loadVector(self: *Self, memref: base.Value, index: base.Value, width: u32) !base.Value {
446 return self.owner().loadVector(memref, index, width);
447 }
448
449 pub fn extractLane(self: *Self, vec: base.Value, lane: u32, dtype: DType) !base.Value {
450 return self.owner().extractLane(vec, lane, dtype);
451 }
452
453 pub fn insertLane(self: *Self, vec: base.Value, scalar: base.Value, lane: u32) !base.Value {
454 return self.owner().insertLane(vec, scalar, lane);
455 }
456
457 pub fn packVector(self: *Self, lanes: [4]base.Value) !base.Value {
458 return self.owner().packVector(lanes);
459 }
460
461 pub fn storeIndex(self: *Self, value: base.Value, memref: base.Value, index: anytype) !void {
462 try self.owner().storeIndex(value, memref, index);
463 }
464
465 pub fn atomicRmw(
466 self: *Self,
467 kind: base.AtomicRmwKind,
468 value: base.Value,
469 memref: base.Value,
470 index: base.Value,
471 ) !base.Value {
472 return self.owner().atomicRmw(kind, value, memref, index);
473 }
474
475 pub fn atomicRmwIndex(
476 self: *Self,
477 kind: base.AtomicRmwKind,
478 value: base.Value,
479 memref: base.Value,
480 index: anytype,
481 ) !base.Value {
482 return self.atomicRmw(kind, value, memref, try self.linearIndex(index));
483 }
484
485 pub fn atomicCas(
486 self: *Self,
487 expected: base.Value,
488 desired: base.Value,
489 memref: base.Value,
490 index: base.Value,
491 ) !base.Value {
492 return self.owner().atomicCas(expected, desired, memref, index);
493 }
494
495 pub fn atomicCasIndex(
496 self: *Self,
497 expected: base.Value,
498 desired: base.Value,
499 memref: base.Value,
500 index: anytype,
501 ) !base.Value {
502 return self.atomicCas(expected, desired, memref, try self.linearIndex(index));
503 }
504
505 pub fn vec2(self: *Self, x: base.Value, y: base.Value) vector.Vec2 {
506 return self.owner().vec2(x, y);
507 }
508
509 pub fn typedVec2(self: *Self, comptime dtype: DType, x: anytype, y: anytype) !typed.Vec2(dtype) {
510 return self.owner().typedVec2(dtype, x, y);
511 }
512
513 pub fn vec3(self: *Self, x: base.Value, y: base.Value, z: base.Value) vector.Vec3 {
514 return self.owner().vec3(x, y, z);
515 }
516
517 pub fn typedVec3(self: *Self, comptime dtype: DType, x: anytype, y: anytype, z: anytype) !typed.Vec3(dtype) {
518 return self.owner().typedVec3(dtype, x, y, z);
519 }
520
521 pub fn splat2(self: *Self, value: base.Value) vector.Vec2 {
522 return self.owner().splat2(value);
523 }
524
525 pub fn typedSplat2(self: *Self, comptime dtype: DType, value: anytype) !typed.Vec2(dtype) {
526 return self.owner().typedSplat2(dtype, value);
527 }
528
529 pub fn splat3(self: *Self, value: base.Value) vector.Vec3 {
530 return self.owner().splat3(value);
531 }
532
533 pub fn typedSplat3(self: *Self, comptime dtype: DType, value: anytype) !typed.Vec3(dtype) {
534 return self.owner().typedSplat3(dtype, value);
535 }
536
537 fn owner(self: *Self) *Inner {
538 return @ptrCast(@alignCast(self.token));
539 }
540 };
541 }
542
543 pub fn wrap(inner: anytype, schedule: anytype) Builder(PointerChild(@TypeOf(inner)), @TypeOf(schedule)) {
544 return Builder(PointerChild(@TypeOf(inner)), @TypeOf(schedule)).init(inner, schedule);
545 }
546
547 fn PointerChild(comptime Pointer: type) type {
548 return switch (@typeInfo(Pointer)) {
549 .pointer => |info| info.child,
550 else => @compileError("logical kernel builders wrap builder pointers"),
551 };
552 }
553
554 fn WhileScope(comptime Inner: type) type {
555 const Return = @TypeOf(@as(*Inner, undefined).whileScope(
556 @as([]const base.Value, undefined),
557 @as([]const base.Type, undefined),
558 ));
559 return switch (@typeInfo(Return)) {
560 .error_union => |info| info.payload,
561 else => @compileError("whileScope must return an error union"),
562 };
563 }
564
565 fn ForScope(comptime Inner: type) type {
566 const Return = @TypeOf(@as(*Inner, undefined).forScope(
567 @as(base.Value, undefined),
568 @as(base.Value, undefined),
569 @as(base.Value, undefined),
570 @as([]const base.Value, undefined),
571 @as([]const base.Type, undefined),
572 ));
573 return switch (@typeInfo(Return)) {
574 .error_union => |info| info.payload,
575 else => @compileError("forScope must return an error union"),
576 };
577 }