lib/choir/src/backends/gpu/nvptx/ptx.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const abi = @import("choir_abi");
3 const choir_pkg = @import("../../../root.zig");
4
5 const nvptx = @import("root.zig");
6 const gpu = @import("../../../dialects/gpu/root.zig");
7
8 const ir = choir_pkg.ir;
9 const dialects = choir_pkg.dialects;
10
11 const Allocator = std.mem.Allocator;
12 const ArithDialect = dialects.ArithDialect;
13 const CmpPredicate = dialects.arith.CmpPredicate;
14 const arith_names = dialects.arith.type_names;
15 const BuiltinDialect = dialects.BuiltinDialect;
16 const FuncDialect = dialects.FuncDialect;
17 const MemrefDialect = dialects.MemrefDialect;
18 const NvptxDialect = nvptx.NvptxDialect;
19 const ScfDialect = dialects.ScfDialect;
20
21 const EmitError = abi.Error || std.Io.Writer.Error;
22
23 const ScalarKind = dialects.arith.ScalarKind;
24 const scalar_kinds = dialects.arith.ScalarSet.init(&.{
25 .i8,
26 .i16,
27 .i32,
28 .i64,
29 .u8,
30 .u16,
31 .u32,
32 .u64,
33 .f16,
34 .bf16,
35 .f32,
36 .f64,
37 .index,
38 .bool,
39 });
40
41 const Value = union(enum) {
42 pred: u32,
43 u32: u32,
44 s32: u32,
45 u64: u32,
46 f32: u32,
47 f32x4: u32,
48 f64: u32,
49 ptr: u32,
50 shared: u32,
51 };
52
53 const AddressBaseKey = struct {
54 space: u32,
55 base: usize,
56 bytes: u32,
57 };
58
59 const AddressParts = struct {
60 reg: u32,
61 imm: i64,
62 };
63
64 const max_address_immediate: i64 = 1 << 30;
65
66 const PeeledIndex = struct {
67 base: ?*ir.Value,
68 offset: i64,
69 };
70
71 fn constantIndexValue(value: *ir.Value) ?i64 {
72 const def = value.getDefiningOp() orelse return null;
73 const op: *ir.Operation = @ptrCast(@alignCast(def));
74 if (!std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) return null;
75 const constant = ArithDialect.ConstantOp{ .op = op };
76 return constant.getIntValue();
77 }
78
79 fn peelIndexOffset(index: *ir.Value) PeeledIndex {
80 var current: *ir.Value = index;
81 var offset: i64 = 0;
82 var depth: u32 = 0;
83 while (depth < 16) : (depth += 1) {
84 if (constantIndexValue(current)) |value| return .{
85 .base = null,
86 .offset = std.math.add(i64, offset, value) catch return .{ .base = index, .offset = 0 },
87 };
88 const def = current.getDefiningOp() orelse break;
89 const op: *ir.Operation = @ptrCast(@alignCast(def));
90 if (std.mem.eql(u8, op.name.name, ArithDialect.AddOp.operation_name)) {
91 const lhs = op.operands.items[0].value;
92 const rhs = op.operands.items[1].value;
93 if (constantIndexValue(rhs)) |value| {
94 offset = std.math.add(i64, offset, value) catch
95 return .{ .base = index, .offset = 0 };
96 current = lhs;
97 continue;
98 }
99 if (constantIndexValue(lhs)) |value| {
100 offset = std.math.add(i64, offset, value) catch
101 return .{ .base = index, .offset = 0 };
102 current = rhs;
103 continue;
104 }
105 break;
106 }
107 if (std.mem.eql(u8, op.name.name, ArithDialect.SubOp.operation_name)) {
108 if (constantIndexValue(op.operands.items[1].value)) |value| {
109 offset -= value;
110 current = op.operands.items[0].value;
111 continue;
112 }
113 break;
114 }
115 break;
116 }
117 return .{ .base = current, .offset = offset };
118 }
119
120 fn intValueForKind(kind: ScalarKind, reg: u32) Value {
121 return switch (kind) {
122 .i8, .i16, .i32 => .{ .s32 = reg },
123 .index, .u8, .u16, .u32 => .{ .u32 = reg },
124 .i64, .u64 => .{ .u64 = reg },
125 else => unreachable,
126 };
127 }
128
129 fn scalarKindIsSignedInteger(kind: ScalarKind) bool {
130 return switch (kind) {
131 .i8, .i16, .i32, .i64 => true,
132 else => false,
133 };
134 }
135
136 const LoopInteger = enum {
137 u32,
138 s32,
139 u64,
140 s64,
141
142 fn fromType(typ: ir.Type) abi.Error!LoopInteger {
143 return switch (try computeKind(typ)) {
144 .index, .u32 => .u32,
145 .i32 => .s32,
146 .u64 => .u64,
147 .i64 => .s64,
148 else => error.UnsupportedOperation,
149 };
150 }
151
152 fn registerFile(self: LoopInteger) []const u8 {
153 return switch (self) {
154 .u32, .s32 => "r",
155 .u64, .s64 => "rd",
156 };
157 }
158
159 fn register(self: LoopInteger, emitter: *Emitter, value: Value) abi.Error!u32 {
160 return switch (self) {
161 .u32, .s32 => emitter.asU32(value),
162 .u64, .s64 => emitter.asU64(value),
163 };
164 }
165 };
166
167 fn atomicIntegerSuffix(kind: dialects.AtomicRmwKind, element: ScalarKind) abi.Error![]const u8 {
168 return switch (kind) {
169 .add => "add.u32",
170 .min => if (element == .i32) "min.s32" else "min.u32",
171 .max => if (element == .i32) "max.s32" else "max.u32",
172 .bit_and => "and.b32",
173 .bit_or => "or.b32",
174 .bit_xor => "xor.b32",
175 .exchange => "exch.b32",
176 };
177 }
178
179 fn atomicKindSupportsRed(kind: dialects.AtomicRmwKind) bool {
180 return switch (kind) {
181 .add, .min, .max, .bit_and, .bit_or, .bit_xor => true,
182 .exchange => false,
183 };
184 }
185
186 pub fn emitPtx(
187 result_allocator: Allocator,
188 entry_name: []const u8,
189 module: *ir.Operation,
190 ) abi.Error![]u8 {
191 var emitter = Emitter.init(result_allocator, entry_name, module);
192 defer emitter.deinit();
193 return emitter.emit() catch |err| switch (err) {
194 error.WriteFailed => error.OutOfMemory,
195 else => |other| other,
196 };
197 }
198
199 const Emitter = struct {
200 allocator: Allocator,
201 entry_name: []const u8,
202 module: *ir.Operation,
203 body: std.Io.Writer.Allocating,
204 memory_decls: std.Io.Writer.Allocating,
205 values: std.AutoHashMapUnmanaged(*const ir.Value, Value) = .{},
206 shared_bases: std.AutoHashMapUnmanaged(AddressBaseKey, u32) = .{},
207 shared_byte_offsets: std.AutoHashMapUnmanaged(u32, u32) = .{},
208 pointer_bases: std.AutoHashMapUnmanaged(AddressBaseKey, u32) = .{},
209 next_r: u32 = 1,
210 next_f: u32 = 1,
211 next_fd: u32 = 1,
212 next_h: u32 = 1,
213 next_rd: u32 = 1,
214 next_p: u32 = 1,
215 next_shared: u32 = 0,
216 next_label: u32 = 0,
217 dynamic_shared_alignment: u32 = 0,
218 requires_sm80: bool = false,
219
220 const OperationHandler = *const fn (*Emitter, *ir.Operation) EmitError!void;
221
222 fn init(allocator: Allocator, entry_name: []const u8, module: *ir.Operation) Emitter {
223 return .{
224 .allocator = allocator,
225 .entry_name = entry_name,
226 .module = module,
227 .body = std.Io.Writer.Allocating.init(allocator),
228 .memory_decls = std.Io.Writer.Allocating.init(allocator),
229 };
230 }
231
232 fn deinit(self: *Emitter) void {
233 self.pointer_bases.deinit(self.allocator);
234 self.shared_byte_offsets.deinit(self.allocator);
235 self.shared_bases.deinit(self.allocator);
236 self.values.deinit(self.allocator);
237 self.memory_decls.deinit();
238 self.body.deinit();
239 }
240
241 fn resetAddressBases(self: *Emitter) void {
242 self.shared_bases.clearRetainingCapacity();
243 self.pointer_bases.clearRetainingCapacity();
244 }
245
246 fn emit(self: *Emitter) EmitError![]u8 {
247 const func = try self.findKernelFunction();
248 try self.emitParameterLoads(func);
249 try self.emitBlock(func.getEntryBlock());
250
251 var out = std.Io.Writer.Allocating.init(self.allocator);
252 errdefer out.deinit();
253 try writeHeader(&out.writer, self.requires_sm80);
254 if (self.dynamic_shared_alignment != 0) {
255 try out.writer.print(" .extern .shared .align {d} .b8 __choir_dynamic_shared[];\n", .{self.dynamic_shared_alignment});
256 try out.writer.writeByte('\n');
257 }
258 try self.emitEntryHeader(&out.writer, func);
259 try self.emitRegisterDecls(&out.writer);
260 try out.writer.writeAll(self.memory_decls.written());
261 if (self.memory_decls.written().len != 0) try out.writer.writeByte('\n');
262 try out.writer.writeAll(self.body.written());
263 try out.writer.writeAll("}\n");
264 return out.toOwnedSlice() catch return error.OutOfMemory;
265 }
266
267 fn findKernelFunction(self: *Emitter) abi.Error!FuncDialect.FuncOp {
268 if (!std.mem.eql(u8, self.module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
269 return error.InvalidArtifact;
270 }
271 const block = self.module.getRegion(0).?.getEntryBlock() orelse return error.InvalidArtifact;
272 var ops = block.getOperations();
273 while (ops.next()) |op| {
274 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) continue;
275 const func = FuncDialect.FuncOp{ .op = op };
276 if (!func.isKernel()) continue;
277 const name = func.getName() orelse return error.InvalidArtifact;
278 if (std.mem.eql(u8, name, self.entry_name)) return func;
279 }
280 return error.InvalidArtifact;
281 }
282
283 fn emitParameterLoads(self: *Emitter, func: FuncDialect.FuncOp) EmitError!void {
284 const args = func.getArguments();
285 for (args, 0..) |arg, index| {
286 if (memrefInfo(arg.type)) |memref| {
287 if (!isKernelParameterAddressSpace(memref.addr_space)) return error.UnsupportedOperation;
288 const raw = self.allocPtr();
289 const ptr = self.allocPtr();
290 try self.line(" ld.param.u64 %rd{d}, [param{d}];", .{ raw, index });
291 try self.line(" cvta.to.global.u64 %rd{d}, %rd{d};", .{ ptr, raw });
292 try self.bind(arg, .{ .ptr = ptr });
293 continue;
294 }
295 switch (try scalarKind(arg.type)) {
296 .i64, .u64 => {
297 const reg = self.allocU64();
298 try self.line(" ld.param.u64 %rd{d}, [param{d}];", .{ reg, index });
299 try self.bind(arg, .{ .u64 = reg });
300 },
301 .index, .u8, .u16, .u32 => {
302 const reg = self.allocU32();
303 try self.line(" ld.param.u32 %r{d}, [param{d}];", .{ reg, index });
304 try self.bind(arg, .{ .u32 = reg });
305 },
306 .i8, .i16, .i32 => {
307 const reg = self.allocU32();
308 try self.line(" ld.param.u32 %r{d}, [param{d}];", .{ reg, index });
309 try self.bind(arg, .{ .s32 = reg });
310 },
311 .f32 => {
312 const reg = self.allocF32();
313 try self.line(" ld.param.f32 %f{d}, [param{d}];", .{ reg, index });
314 try self.bind(arg, .{ .f32 = reg });
315 },
316 .f64 => {
317 const reg = self.allocF64();
318 try self.line(" ld.param.f64 %fd{d}, [param{d}];", .{ reg, index });
319 try self.bind(arg, .{ .f64 = reg });
320 },
321 .bool, .f16, .bf16 => return error.UnsupportedOperation,
322 }
323 }
324 if (args.len != 0) try self.body.writer.writeByte('\n');
325 }
326
327 fn emitEntryHeader(self: *Emitter, writer: *std.Io.Writer, func: FuncDialect.FuncOp) EmitError!void {
328 try writer.print(".visible .entry {s}(\n", .{self.entry_name});
329 const args = func.getArguments();
330 for (args, 0..) |arg, index| {
331 const suffix = if (index + 1 == args.len) "" else ",";
332 if (memrefInfo(arg.type)) |memref| {
333 if (!isKernelParameterAddressSpace(memref.addr_space)) return error.UnsupportedOperation;
334 try writer.print(" .param .u64 param{d}{s}\n", .{ index, suffix });
335 continue;
336 }
337 switch (try scalarKind(arg.type)) {
338 .i64, .u64 => try writer.print(" .param .u64 param{d}{s}\n", .{ index, suffix }),
339 .index, .i8, .i16, .i32, .u8, .u16, .u32 => try writer.print(" .param .u32 param{d}{s}\n", .{ index, suffix }),
340 .f32 => try writer.print(" .param .f32 param{d}{s}\n", .{ index, suffix }),
341 .f64 => try writer.print(" .param .f64 param{d}{s}\n", .{ index, suffix }),
342 .bool, .f16, .bf16 => return error.UnsupportedOperation,
343 }
344 }
345 try writer.writeAll(")\n{\n");
346 }
347
348 fn emitRegisterDecls(self: *Emitter, writer: *std.Io.Writer) EmitError!void {
349 try writer.print(" .reg .b32 %r<{d}>;\n", .{@max(self.next_r, 1)});
350 try writer.print(" .reg .f32 %f<{d}>;\n", .{@max(self.next_f, 1)});
351 try writer.print(" .reg .f64 %fd<{d}>;\n", .{@max(self.next_fd, 1)});
352 try writer.print(" .reg .b16 %h<{d}>;\n", .{@max(self.next_h, 1)});
353 try writer.print(" .reg .b64 %rd<{d}>;\n", .{@max(self.next_rd, 1)});
354 try writer.print(" .reg .pred %p<{d}>;\n\n", .{@max(self.next_p, 1)});
355 }
356
357 fn emitBlock(self: *Emitter, block: *ir.Block) EmitError!void {
358 self.resetAddressBases();
359 var ops = block.getOperations();
360 while (ops.next()) |op| {
361 try self.emitOperation(op);
362 }
363 self.resetAddressBases();
364 }
365
366 fn emitOperation(self: *Emitter, op: *ir.Operation) EmitError!void {
367 const handler = operation_emitters.get(op.name.name) orelse return error.UnsupportedOperation;
368 try handler(self, op);
369 for (op.results.items) |*result| try self.narrowResult(result);
370 }
371
372 fn narrowResult(self: *Emitter, result: *ir.Value) EmitError!void {
373 const kind = scalar_kinds.kindFromType(result.type) orelse return;
374 const bits: u5 = switch (kind) {
375 .i8, .u8 => 8,
376 .i16, .u16 => 16,
377 else => return,
378 };
379 const input = try self.asU32(try self.require(result));
380 const out = self.allocU32();
381 if (kind == .u8 or kind == .u16) {
382 const mask = (@as(u32, 1) << bits) - 1;
383 try self.line(" and.b32 %r{d}, %r{d}, {d};", .{ out, input, mask });
384 } else {
385 const shift: u6 = 32 - @as(u6, bits);
386 try self.line(" shl.b32 %r{d}, %r{d}, {d};", .{ out, input, shift });
387 try self.line(" shr.s32 %r{d}, %r{d}, {d};", .{ out, out, shift });
388 }
389 try self.bind(result, intValueForKind(kind, out));
390 }
391
392 const operation_emitters = std.StaticStringMap(OperationHandler).initComptime(.{
393 .{ FuncDialect.ReturnOp.operation_name, lineHandler(" ret;") },
394 .{ NvptxDialect.ThreadIdxOp.operation_name, dimHandler(NvptxDialect.ThreadIdxOp, "tid") },
395 .{ NvptxDialect.BlockIdxOp.operation_name, dimHandler(NvptxDialect.BlockIdxOp, "ctaid") },
396 .{ NvptxDialect.BlockDimOp.operation_name, dimHandler(NvptxDialect.BlockDimOp, "ntid") },
397 .{ NvptxDialect.GridDimOp.operation_name, dimHandler(NvptxDialect.GridDimOp, "nctaid") },
398 .{ NvptxDialect.Barrier0Op.operation_name, lineHandler(" bar.sync 0;") },
399 .{ NvptxDialect.WarpBarrierAllOp.operation_name, lineHandler(" bar.warp.sync 0xffffffff;") },
400 .{ NvptxDialect.LoadGlobalOp.operation_name, wrappedHandler(NvptxDialect.LoadGlobalOp, "emitNvptxGlobalLoad") },
401 .{ NvptxDialect.LoadLocalOp.operation_name, wrappedHandler(NvptxDialect.LoadLocalOp, "emitNvptxLocalLoad") },
402 .{ NvptxDialect.LoadSharedOp.operation_name, wrappedHandler(NvptxDialect.LoadSharedOp, "emitNvptxSharedLoad") },
403 .{ NvptxDialect.StoreGlobalOp.operation_name, wrappedHandler(NvptxDialect.StoreGlobalOp, "emitNvptxGlobalStore") },
404 .{ NvptxDialect.StoreLocalOp.operation_name, wrappedHandler(NvptxDialect.StoreLocalOp, "emitNvptxLocalStore") },
405 .{ NvptxDialect.StoreSharedOp.operation_name, wrappedHandler(NvptxDialect.StoreSharedOp, "emitNvptxSharedStore") },
406 .{ NvptxDialect.AtomicGlobalOp.operation_name, wrappedHandler(NvptxDialect.AtomicGlobalOp, "emitNvptxGlobalAtomic") },
407 .{ NvptxDialect.AtomicSharedOp.operation_name, wrappedHandler(NvptxDialect.AtomicSharedOp, "emitNvptxSharedAtomic") },
408 .{ NvptxDialect.AtomicCasGlobalOp.operation_name, wrappedHandler(NvptxDialect.AtomicCasGlobalOp, "emitNvptxGlobalAtomicCas") },
409 .{ NvptxDialect.AtomicCasSharedOp.operation_name, wrappedHandler(NvptxDialect.AtomicCasSharedOp, "emitNvptxSharedAtomicCas") },
410 .{ NvptxDialect.LaneIdOp.operation_name, specialHandler("laneid") },
411 .{ NvptxDialect.WarpIdOp.operation_name, rawHandler("emitWarpIndex") },
412 .{ NvptxDialect.SyncWarpOp.operation_name, wrappedHandler(NvptxDialect.SyncWarpOp, "emitSyncWarp") },
413 .{ NvptxDialect.ActiveMaskOp.operation_name, wrappedHandler(NvptxDialect.ActiveMaskOp, "emitActiveMask") },
414 .{ NvptxDialect.AllSyncOp.operation_name, voteHandler(NvptxDialect.AllSyncOp, "all") },
415 .{ NvptxDialect.AnySyncOp.operation_name, voteHandler(NvptxDialect.AnySyncOp, "any") },
416 .{ NvptxDialect.BallotSyncOp.operation_name, wrappedHandler(NvptxDialect.BallotSyncOp, "emitBallotSync") },
417 .{ NvptxDialect.ShflSyncOp.operation_name, wrappedHandler(NvptxDialect.ShflSyncOp, "emitShflSync") },
418 .{ NvptxDialect.WarpReduceOp.operation_name, wrappedHandler(NvptxDialect.WarpReduceOp, "emitWarpReduce") },
419 .{ NvptxDialect.WarpScanOp.operation_name, wrappedHandler(NvptxDialect.WarpScanOp, "emitWarpScan") },
420 .{ NvptxDialect.MmaSyncOp.operation_name, wrappedHandler(NvptxDialect.MmaSyncOp, "emitMmaSync") },
421 .{ NvptxDialect.FenceDeviceOp.operation_name, lineHandler(" membar.gl;") },
422 .{ NvptxDialect.CpAsyncSharedOp.operation_name, wrappedHandler(NvptxDialect.CpAsyncSharedOp, "emitCpAsyncShared") },
423 .{ NvptxDialect.CpAsyncCommitOp.operation_name, wrappedHandler(NvptxDialect.CpAsyncCommitOp, "emitCpAsyncCommit") },
424 .{ NvptxDialect.CpAsyncWaitOp.operation_name, wrappedHandler(NvptxDialect.CpAsyncWaitOp, "emitCpAsyncWait") },
425 .{ MemrefDialect.AllocOp.operation_name, wrappedHandler(MemrefDialect.AllocOp, "emitAlloc") },
426 .{ MemrefDialect.AllocaOp.operation_name, wrappedHandler(MemrefDialect.AllocaOp, "emitAlloca") },
427 .{ ArithDialect.ConstantOp.operation_name, wrappedHandler(ArithDialect.ConstantOp, "emitConstant") },
428 .{ ArithDialect.AddOp.operation_name, binaryHandler("add") },
429 .{ ArithDialect.SubOp.operation_name, binaryHandler("sub") },
430 .{ ArithDialect.MulOp.operation_name, binaryHandler("mul") },
431 .{ ArithDialect.UmulhiOp.operation_name, rawHandler("emitUmulhi") },
432 .{ ArithDialect.DivOp.operation_name, rawHandler("emitDiv") },
433 .{ ArithDialect.NegOp.operation_name, rawHandler("emitNeg") },
434 .{ ArithDialect.AbsOp.operation_name, rawHandler("emitAbs") },
435 .{ ArithDialect.SqrtOp.operation_name, floatUnaryHandler("sqrt.approx") },
436 .{ ArithDialect.MaxOp.operation_name, binaryHandler("max") },
437 .{ ArithDialect.MinOp.operation_name, binaryHandler("min") },
438 .{ ArithDialect.AndOp.operation_name, bitwiseHandler("and") },
439 .{ ArithDialect.OrOp.operation_name, bitwiseHandler("or") },
440 .{ ArithDialect.XorOp.operation_name, bitwiseHandler("xor") },
441 .{ ArithDialect.NotOp.operation_name, rawHandler("emitNot") },
442 .{ ArithDialect.PopCountOp.operation_name, rawHandler("emitPopCount") },
443 .{ ArithDialect.ShlOp.operation_name, shiftHandler("shl.b32") },
444 .{ ArithDialect.ShrOp.operation_name, shiftHandler("shr.s32") },
445 .{ ArithDialect.UshrOp.operation_name, shiftHandler("shr.u32") },
446 .{ ArithDialect.CmpOp.operation_name, wrappedHandler(ArithDialect.CmpOp, "emitCompare") },
447 .{ ArithDialect.SelectOp.operation_name, rawHandler("emitSelect") },
448 .{ ArithDialect.CastOp.operation_name, wrappedHandler(ArithDialect.CastOp, "emitCast") },
449 .{ ArithDialect.BitcastOp.operation_name, rawHandler("emitBitcast") },
450 .{ ArithDialect.ExpOp.operation_name, rawHandler("emitExp") },
451 .{ ArithDialect.LogOp.operation_name, rawHandler("emitLog") },
452 .{ ArithDialect.TanhOp.operation_name, rawHandler("emitTanh") },
453 .{ ArithDialect.SinOp.operation_name, floatUnaryHandler("sin.approx") },
454 .{ ArithDialect.CosOp.operation_name, floatUnaryHandler("cos.approx") },
455 .{ ArithDialect.TanOp.operation_name, rawHandler("emitTan") },
456 .{ ArithDialect.FloorOp.operation_name, floatIntegralUnaryHandler("cvt.rmi") },
457 .{ ArithDialect.RoundOp.operation_name, rawHandler("emitRound") },
458 .{ ArithDialect.TruncOp.operation_name, floatIntegralUnaryHandler("cvt.rzi") },
459 .{ ArithDialect.Tf32RoundOp.operation_name, rawHandler("emitTf32Round") },
460 .{ ArithDialect.PowOp.operation_name, rawHandler("emitPow") },
461 .{ ArithDialect.Atan2Op.operation_name, rawHandler("emitAtan2") },
462 .{ ArithDialect.FmaOp.operation_name, rawHandler("emitFma") },
463 .{ ArithDialect.ExtractOp.operation_name, rawHandler("emitVecExtract") },
464 .{ ArithDialect.SplatOp.operation_name, rawHandler("emitVecSplat") },
465 .{ ArithDialect.InsertOp.operation_name, rawHandler("emitVecInsert") },
466 .{ ScfDialect.IfOp.operation_name, wrappedHandler(ScfDialect.IfOp, "emitIf") },
467 .{ ScfDialect.ForOp.operation_name, wrappedHandler(ScfDialect.ForOp, "emitFor") },
468 .{ ScfDialect.WhileOp.operation_name, wrappedHandler(ScfDialect.WhileOp, "emitWhile") },
469 .{ ScfDialect.YieldOp.operation_name, nopHandler() },
470 });
471
472 fn lineHandler(comptime text: []const u8) OperationHandler {
473 return struct {
474 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
475 _ = op;
476 try self.line(text, .{});
477 }
478 }.emit;
479 }
480
481 fn nopHandler() OperationHandler {
482 return struct {
483 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
484 _ = self;
485 _ = op;
486 }
487 }.emit;
488 }
489
490 fn dimHandler(comptime OpType: type, comptime register: []const u8) OperationHandler {
491 return struct {
492 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
493 const wrapped_op = OpType{ .op = op };
494 try self.emitDimRegister(op, wrapped_op.getDimension() orelse return error.InvalidArtifact, register);
495 }
496 }.emit;
497 }
498
499 fn specialHandler(comptime register: []const u8) OperationHandler {
500 return struct {
501 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
502 try self.emitSpecialRegister(op, register);
503 }
504 }.emit;
505 }
506
507 fn voteHandler(comptime OpType: type, comptime kind: []const u8) OperationHandler {
508 return struct {
509 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
510 const wrapped_op = OpType{ .op = op };
511 try self.emitVoteSync(wrapped_op.getResult(), wrapped_op.getMask(), wrapped_op.getPredicate(), kind);
512 }
513 }.emit;
514 }
515
516 fn wrappedHandler(comptime OpType: type, comptime method: []const u8) OperationHandler {
517 return struct {
518 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
519 try @field(Emitter, method)(self, OpType{ .op = op });
520 }
521 }.emit;
522 }
523
524 fn rawHandler(comptime method: []const u8) OperationHandler {
525 return struct {
526 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
527 try @field(Emitter, method)(self, op);
528 }
529 }.emit;
530 }
531
532 fn binaryHandler(comptime mnemonic: []const u8) OperationHandler {
533 return struct {
534 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
535 try self.emitBinary(op, mnemonic);
536 }
537 }.emit;
538 }
539
540 fn bitwiseHandler(comptime mnemonic: []const u8) OperationHandler {
541 return struct {
542 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
543 try self.emitBitwise(op, mnemonic);
544 }
545 }.emit;
546 }
547
548 fn shiftHandler(comptime mnemonic: []const u8) OperationHandler {
549 return struct {
550 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
551 try self.emitShift(op, mnemonic);
552 }
553 }.emit;
554 }
555
556 fn floatUnaryHandler(comptime mnemonic: []const u8) OperationHandler {
557 return struct {
558 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
559 try self.emitFloatUnary(op, mnemonic);
560 }
561 }.emit;
562 }
563
564 fn floatIntegralUnaryHandler(comptime mnemonic: []const u8) OperationHandler {
565 return struct {
566 fn emit(self: *Emitter, op: *ir.Operation) EmitError!void {
567 try self.emitFloatIntegralUnary(op, mnemonic);
568 }
569 }.emit;
570 }
571 fn emitDimRegister(self: *Emitter, op: *ir.Operation, dim: gpu.Dimension, comptime ptx_name: []const u8) EmitError!void {
572 const out = self.allocU32();
573 try self.line(" mov.u32 %r{d}, %{s}.{s};", .{ out, ptx_name, dimName(dim) });
574 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .u32 = out });
575 }
576
577 fn emitSpecialRegister(self: *Emitter, op: *ir.Operation, comptime ptx_name: []const u8) EmitError!void {
578 const out = self.allocU32();
579 try self.line(" mov.u32 %r{d}, %{s};", .{ out, ptx_name });
580 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .u32 = out });
581 }
582
583 fn emitWarpIndex(self: *Emitter, op: *ir.Operation) EmitError!void {
584 const tid_z = self.allocU32();
585 try self.line(" mov.u32 %r{d}, %tid.z;", .{tid_z});
586 const ntid_y = self.allocU32();
587 try self.line(" mov.u32 %r{d}, %ntid.y;", .{ntid_y});
588 const tid_y = self.allocU32();
589 try self.line(" mov.u32 %r{d}, %tid.y;", .{tid_y});
590 const ntid_x = self.allocU32();
591 try self.line(" mov.u32 %r{d}, %ntid.x;", .{ntid_x});
592 const tid_x = self.allocU32();
593 try self.line(" mov.u32 %r{d}, %tid.x;", .{tid_x});
594 const plane = self.allocU32();
595 try self.line(" mul.lo.u32 %r{d}, %r{d}, %r{d};", .{ plane, tid_z, ntid_y });
596 const row = self.allocU32();
597 try self.line(" add.u32 %r{d}, %r{d}, %r{d};", .{ row, plane, tid_y });
598 const scaled = self.allocU32();
599 try self.line(" mul.lo.u32 %r{d}, %r{d}, %r{d};", .{ scaled, row, ntid_x });
600 const linear = self.allocU32();
601 try self.line(" add.u32 %r{d}, %r{d}, %r{d};", .{ linear, scaled, tid_x });
602 const out = self.allocU32();
603 try self.line(" shr.u32 %r{d}, %r{d}, 5;", .{ out, linear });
604 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .u32 = out });
605 }
606
607 fn emitSyncWarp(self: *Emitter, op: NvptxDialect.SyncWarpOp) EmitError!void {
608 const mask = try self.asU32(try self.require(op.getMask()));
609 try self.line(" bar.warp.sync %r{d};", .{mask});
610 }
611
612 fn emitActiveMask(self: *Emitter, op: NvptxDialect.ActiveMaskOp) EmitError!void {
613 const out = self.allocU32();
614 try self.line(" activemask.b32 %r{d};", .{out});
615 try self.bind(op.getResult(), .{ .u32 = out });
616 }
617
618 fn emitVoteSync(
619 self: *Emitter,
620 result: *ir.Value,
621 mask_value: *ir.Value,
622 predicate_value: *ir.Value,
623 comptime mode: []const u8,
624 ) EmitError!void {
625 if ((try scalarKind(result.type)) != .bool) return error.UnsupportedOperation;
626 const mask = try self.asU32(try self.require(mask_value));
627 const predicate = try self.asPred(try self.require(predicate_value));
628 const out = self.allocPred();
629 try self.line(" vote.sync.{s}.pred %p{d}, %p{d}, %r{d};", .{ mode, out, predicate, mask });
630 try self.bind(result, .{ .pred = out });
631 }
632
633 fn emitBallotSync(self: *Emitter, op: NvptxDialect.BallotSyncOp) EmitError!void {
634 const mask = try self.asU32(try self.require(op.getMask()));
635 const predicate = try self.asPred(try self.require(op.getPredicate()));
636 const out = self.allocU32();
637 try self.line(" vote.sync.ballot.b32 %r{d}, %p{d}, %r{d};", .{ out, predicate, mask });
638 try self.bind(op.getResult(), .{ .u32 = out });
639 }
640
641 fn emitShflSync(self: *Emitter, op: NvptxDialect.ShflSyncOp) EmitError!void {
642 const result = op.getResult();
643 const mask = try self.asU32(try self.require(op.getMask()));
644 const lane_or_delta = try self.asU32(try self.require(op.getLaneOrDelta()));
645 const mode = ptxShuffleMode(op.getMode() orelse return error.InvalidArtifact);
646 const src = try self.require(op.getSrc());
647 switch (try scalarKind(result.type)) {
648 .i64, .u64 => return error.UnsupportedOperation,
649 .index, .i8, .i16, .i32, .u8, .u16, .u32 => {
650 const out = self.allocU32();
651 try self.line(" shfl.sync.{s}.b32 %r{d}, %r{d}, %r{d}, 0x1f, %r{d};", .{ mode, out, try self.asU32(src), lane_or_delta, mask });
652 try self.bind(result, intValueForKind(try scalarKind(result.type), out));
653 },
654 .f32 => {
655 const src_bits = self.allocU32();
656 const out_bits = self.allocU32();
657 const out = self.allocF32();
658 try self.line(" mov.b32 %r{d}, %f{d};", .{ src_bits, try self.asF32(src) });
659 try self.line(" shfl.sync.{s}.b32 %r{d}, %r{d}, %r{d}, 0x1f, %r{d};", .{ mode, out_bits, src_bits, lane_or_delta, mask });
660 try self.line(" mov.b32 %f{d}, %r{d};", .{ out, out_bits });
661 try self.bind(result, .{ .f32 = out });
662 },
663 .f64 => {
664 const src_lo = self.allocU32();
665 const src_hi = self.allocU32();
666 const out_lo = self.allocU32();
667 const out_hi = self.allocU32();
668 const out = self.allocF64();
669 try self.emitF64Unpack(src_lo, src_hi, try self.asF64(src));
670 try self.line(" shfl.sync.{s}.b32 %r{d}, %r{d}, %r{d}, 0x1f, %r{d};", .{ mode, out_lo, src_lo, lane_or_delta, mask });
671 try self.line(" shfl.sync.{s}.b32 %r{d}, %r{d}, %r{d}, 0x1f, %r{d};", .{ mode, out_hi, src_hi, lane_or_delta, mask });
672 try self.emitF64Pack(out, out_lo, out_hi);
673 try self.bind(result, .{ .f64 = out });
674 },
675 .f16, .bf16 => return error.UnsupportedOperation,
676 .bool => return error.UnsupportedOperation,
677 }
678 }
679
680 fn emitWarpReduce(self: *Emitter, op: NvptxDialect.WarpReduceOp) EmitError!void {
681 const result = op.getResult();
682 try self.requireFullWarpMask(op.getMask());
683 const mask = try self.asU32(try self.require(op.getMask()));
684 const value = try self.require(op.getValue());
685 const op_kind = op.getOpKind() orelse return error.InvalidArtifact;
686 switch (try scalarKind(result.type)) {
687 .i64, .u64 => return error.UnsupportedOperation,
688 .index, .i8, .i16, .i32, .u8, .u16, .u32 => {
689 const accumulator = self.allocU32();
690 try self.emitMove(.{ .u32 = accumulator }, value);
691 const unsigned = !scalarKindIsSignedInteger(try scalarKind(result.type));
692 inline for (warp_reduce_deltas) |delta| {
693 const other = self.allocU32();
694 try self.line(" shfl.sync.bfly.b32 %r{d}, %r{d}, {d}, 0x1f, %r{d};", .{ other, accumulator, delta, mask });
695 try self.emitWarpIntegerCombine(null, accumulator, other, op_kind, unsigned);
696 }
697 try self.bind(result, intValueForKind(try scalarKind(result.type), accumulator));
698 },
699 .f32 => {
700 const accumulator = self.allocF32();
701 try self.emitMove(.{ .f32 = accumulator }, value);
702 inline for (warp_reduce_deltas) |delta| {
703 const bits = self.allocU32();
704 const other_bits = self.allocU32();
705 const other = self.allocF32();
706 try self.line(" mov.b32 %r{d}, %f{d};", .{ bits, accumulator });
707 try self.line(" shfl.sync.bfly.b32 %r{d}, %r{d}, {d}, 0x1f, %r{d};", .{ other_bits, bits, delta, mask });
708 try self.line(" mov.b32 %f{d}, %r{d};", .{ other, other_bits });
709 try self.emitWarpFloatCombine(null, accumulator, other, op_kind);
710 }
711 try self.bind(result, .{ .f32 = accumulator });
712 },
713 .f64 => {
714 const accumulator = self.allocF64();
715 try self.emitMove(.{ .f64 = accumulator }, value);
716 inline for (warp_reduce_deltas) |delta| {
717 const bits_lo = self.allocU32();
718 const bits_hi = self.allocU32();
719 const other_lo = self.allocU32();
720 const other_hi = self.allocU32();
721 const other = self.allocF64();
722 try self.emitF64Unpack(bits_lo, bits_hi, accumulator);
723 try self.line(" shfl.sync.bfly.b32 %r{d}, %r{d}, {d}, 0x1f, %r{d};", .{ other_lo, bits_lo, delta, mask });
724 try self.line(" shfl.sync.bfly.b32 %r{d}, %r{d}, {d}, 0x1f, %r{d};", .{ other_hi, bits_hi, delta, mask });
725 try self.emitF64Pack(other, other_lo, other_hi);
726 try self.emitWarpF64Combine(null, accumulator, other, op_kind);
727 }
728 try self.bind(result, .{ .f64 = accumulator });
729 },
730 .f16, .bf16 => return error.UnsupportedOperation,
731 .bool => return error.UnsupportedOperation,
732 }
733 }
734
735 fn emitWarpScan(self: *Emitter, op: NvptxDialect.WarpScanOp) EmitError!void {
736 const result = op.getResult();
737 try self.requireFullWarpMask(op.getMask());
738 const mask = try self.asU32(try self.require(op.getMask()));
739 const value = try self.require(op.getValue());
740 const op_kind = op.getOpKind() orelse return error.InvalidArtifact;
741 switch (try scalarKind(result.type)) {
742 .i64, .u64 => return error.UnsupportedOperation,
743 .index, .i8, .i16, .i32, .u8, .u16, .u32 => try self.emitWarpIntegerScan(result, mask, value, op_kind, op.isInclusive()),
744 .f32 => try self.emitWarpFloatScan(result, mask, value, op_kind, op.isInclusive()),
745 .f16, .bf16 => return error.UnsupportedOperation,
746 .f64 => return error.UnsupportedOperation,
747 .bool => return error.UnsupportedOperation,
748 }
749 }
750
751 fn emitWarpIntegerScan(
752 self: *Emitter,
753 result: *ir.Value,
754 mask: u32,
755 value: Value,
756 op_kind: gpu.WarpOpKind,
757 inclusive: bool,
758 ) EmitError!void {
759 const accumulator = self.allocU32();
760 try self.emitMove(.{ .u32 = accumulator }, value);
761 const unsigned = !scalarKindIsSignedInteger(try scalarKind(result.type));
762 inline for (warp_scan_offsets) |offset| {
763 const other = self.allocU32();
764 const valid = self.allocPred();
765 try self.line(" shfl.sync.up.b32 %r{d}|%p{d}, %r{d}, {d}, 0x0, %r{d};", .{ other, valid, accumulator, offset, mask });
766 try self.emitWarpIntegerCombine(valid, accumulator, other, op_kind, unsigned);
767 }
768 if (inclusive) {
769 try self.bind(result, intValueForKind(try scalarKind(result.type), accumulator));
770 return;
771 }
772
773 const previous = self.allocU32();
774 const previous_valid = self.allocPred();
775 const out = self.allocU32();
776 try self.line(" shfl.sync.up.b32 %r{d}|%p{d}, %r{d}, 1, 0x0, %r{d};", .{ previous, previous_valid, accumulator, mask });
777 try self.emitWarpIntegerIdentity(out, op_kind, unsigned);
778 try self.line(" @%p{d} mov.u32 %r{d}, %r{d};", .{ previous_valid, out, previous });
779 try self.bind(result, intValueForKind(try scalarKind(result.type), out));
780 }
781
782 fn emitWarpFloatScan(
783 self: *Emitter,
784 result: *ir.Value,
785 mask: u32,
786 value: Value,
787 op_kind: gpu.WarpOpKind,
788 inclusive: bool,
789 ) EmitError!void {
790 const accumulator = self.allocF32();
791 try self.emitMove(.{ .f32 = accumulator }, value);
792 inline for (warp_scan_offsets) |offset| {
793 const bits = self.allocU32();
794 const other_bits = self.allocU32();
795 const other = self.allocF32();
796 const valid = self.allocPred();
797 try self.line(" mov.b32 %r{d}, %f{d};", .{ bits, accumulator });
798 try self.line(" shfl.sync.up.b32 %r{d}|%p{d}, %r{d}, {d}, 0x0, %r{d};", .{ other_bits, valid, bits, offset, mask });
799 try self.line(" mov.b32 %f{d}, %r{d};", .{ other, other_bits });
800 try self.emitWarpFloatCombine(valid, accumulator, other, op_kind);
801 }
802 if (inclusive) {
803 try self.bind(result, .{ .f32 = accumulator });
804 return;
805 }
806
807 const bits = self.allocU32();
808 const previous_bits = self.allocU32();
809 const previous = self.allocF32();
810 const previous_valid = self.allocPred();
811 const out = self.allocF32();
812 try self.line(" mov.b32 %r{d}, %f{d};", .{ bits, accumulator });
813 try self.line(" shfl.sync.up.b32 %r{d}|%p{d}, %r{d}, 1, 0x0, %r{d};", .{ previous_bits, previous_valid, bits, mask });
814 try self.line(" mov.b32 %f{d}, %r{d};", .{ previous, previous_bits });
815 try self.emitWarpFloatIdentity(out, op_kind);
816 try self.line(" @%p{d} mov.f32 %f{d}, %f{d};", .{ previous_valid, out, previous });
817 try self.bind(result, .{ .f32 = out });
818 }
819
820 fn emitWarpIntegerCombine(self: *Emitter, predicate: ?u32, accumulator: u32, other: u32, op_kind: gpu.WarpOpKind, unsigned: bool) EmitError!void {
821 if (predicate) |pred| {
822 switch (op_kind) {
823 .add => try self.line(" @%p{d} add.u32 %r{d}, %r{d}, %r{d};", .{ pred, accumulator, accumulator, other }),
824 .max => try self.line(" @%p{d} max.{s}32 %r{d}, %r{d}, %r{d};", .{ pred, if (unsigned) "u" else "s", accumulator, accumulator, other }),
825 .min => try self.line(" @%p{d} min.{s}32 %r{d}, %r{d}, %r{d};", .{ pred, if (unsigned) "u" else "s", accumulator, accumulator, other }),
826 .and_ => try self.line(" @%p{d} and.b32 %r{d}, %r{d}, %r{d};", .{ pred, accumulator, accumulator, other }),
827 .or_ => try self.line(" @%p{d} or.b32 %r{d}, %r{d}, %r{d};", .{ pred, accumulator, accumulator, other }),
828 .xor => try self.line(" @%p{d} xor.b32 %r{d}, %r{d}, %r{d};", .{ pred, accumulator, accumulator, other }),
829 }
830 return;
831 }
832 switch (op_kind) {
833 .add => try self.line(" add.u32 %r{d}, %r{d}, %r{d};", .{ accumulator, accumulator, other }),
834 .max => try self.line(" max.{s}32 %r{d}, %r{d}, %r{d};", .{ if (unsigned) "u" else "s", accumulator, accumulator, other }),
835 .min => try self.line(" min.{s}32 %r{d}, %r{d}, %r{d};", .{ if (unsigned) "u" else "s", accumulator, accumulator, other }),
836 .and_ => try self.line(" and.b32 %r{d}, %r{d}, %r{d};", .{ accumulator, accumulator, other }),
837 .or_ => try self.line(" or.b32 %r{d}, %r{d}, %r{d};", .{ accumulator, accumulator, other }),
838 .xor => try self.line(" xor.b32 %r{d}, %r{d}, %r{d};", .{ accumulator, accumulator, other }),
839 }
840 }
841
842 fn emitWarpFloatCombine(self: *Emitter, predicate: ?u32, accumulator: u32, other: u32, op_kind: gpu.WarpOpKind) EmitError!void {
843 if (predicate) |pred| {
844 switch (op_kind) {
845 .add => try self.line(" @%p{d} add.f32 %f{d}, %f{d}, %f{d};", .{ pred, accumulator, accumulator, other }),
846 .max => try self.line(" @%p{d} max.f32 %f{d}, %f{d}, %f{d};", .{ pred, accumulator, accumulator, other }),
847 .min => try self.line(" @%p{d} min.f32 %f{d}, %f{d}, %f{d};", .{ pred, accumulator, accumulator, other }),
848 .and_, .or_, .xor => return error.UnsupportedOperation,
849 }
850 return;
851 }
852 switch (op_kind) {
853 .add => try self.line(" add.f32 %f{d}, %f{d}, %f{d};", .{ accumulator, accumulator, other }),
854 .max => try self.line(" max.f32 %f{d}, %f{d}, %f{d};", .{ accumulator, accumulator, other }),
855 .min => try self.line(" min.f32 %f{d}, %f{d}, %f{d};", .{ accumulator, accumulator, other }),
856 .and_, .or_, .xor => return error.UnsupportedOperation,
857 }
858 }
859
860 fn emitWarpF64Combine(self: *Emitter, predicate: ?u32, accumulator: u32, other: u32, op_kind: gpu.WarpOpKind) EmitError!void {
861 if (predicate) |pred| {
862 switch (op_kind) {
863 .add => try self.line(" @%p{d} add.f64 %fd{d}, %fd{d}, %fd{d};", .{ pred, accumulator, accumulator, other }),
864 .max => try self.line(" @%p{d} max.f64 %fd{d}, %fd{d}, %fd{d};", .{ pred, accumulator, accumulator, other }),
865 .min => try self.line(" @%p{d} min.f64 %fd{d}, %fd{d}, %fd{d};", .{ pred, accumulator, accumulator, other }),
866 .and_, .or_, .xor => return error.UnsupportedOperation,
867 }
868 return;
869 }
870 switch (op_kind) {
871 .add => try self.line(" add.f64 %fd{d}, %fd{d}, %fd{d};", .{ accumulator, accumulator, other }),
872 .max => try self.line(" max.f64 %fd{d}, %fd{d}, %fd{d};", .{ accumulator, accumulator, other }),
873 .min => try self.line(" min.f64 %fd{d}, %fd{d}, %fd{d};", .{ accumulator, accumulator, other }),
874 .and_, .or_, .xor => return error.UnsupportedOperation,
875 }
876 }
877
878 fn emitF64Unpack(self: *Emitter, lo: u32, hi: u32, src: u32) EmitError!void {
879 try self.line(" mov.b64 {{%r{d}, %r{d}}}, %fd{d};", .{ lo, hi, src });
880 }
881
882 fn emitF64Pack(self: *Emitter, dst: u32, lo: u32, hi: u32) EmitError!void {
883 try self.line(" mov.b64 %fd{d}, {{%r{d}, %r{d}}};", .{ dst, lo, hi });
884 }
885
886 fn emitWarpIntegerIdentity(self: *Emitter, destination: u32, op_kind: gpu.WarpOpKind, unsigned: bool) EmitError!void {
887 switch (op_kind) {
888 .add, .or_, .xor => try self.line(" mov.u32 %r{d}, 0;", .{destination}),
889 .and_ => try self.line(" mov.u32 %r{d}, 4294967295;", .{destination}),
890 .max => try self.line(" mov.u32 %r{d}, {d};", .{ destination, if (unsigned) @as(u32, 0) else @as(u32, 2147483648) }),
891 .min => try self.line(" mov.u32 %r{d}, {d};", .{ destination, if (unsigned) @as(u32, 4294967295) else @as(u32, 2147483647) }),
892 }
893 }
894
895 fn emitWarpFloatIdentity(self: *Emitter, destination: u32, op_kind: gpu.WarpOpKind) EmitError!void {
896 switch (op_kind) {
897 .add => try self.line(" mov.f32 %f{d}, 0f00000000;", .{destination}),
898 .max => try self.line(" mov.f32 %f{d}, 0fFF800000;", .{destination}),
899 .min => try self.line(" mov.f32 %f{d}, 0f7F800000;", .{destination}),
900 .and_, .or_, .xor => return error.UnsupportedOperation,
901 }
902 }
903
904 fn requireFullWarpMask(_: *Emitter, mask: *ir.Value) abi.Error!void {
905 const defining = mask.getDefiningOp() orelse return error.UnsupportedOperation;
906 const op: *ir.Operation = @ptrCast(@alignCast(defining));
907 if (!std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) return error.UnsupportedOperation;
908 const constant = ArithDialect.ConstantOp{ .op = op };
909 const int_value = constant.getIntValue() orelse return error.UnsupportedOperation;
910 if (int_value != -1) return error.UnsupportedOperation;
911 }
912
913 fn emitAlloc(self: *Emitter, op: MemrefDialect.AllocOp) EmitError!void {
914 const result = op.getResult();
915 const memref = memrefInfo(result.type) orelse return error.InvalidArtifact;
916 if (memref.addr_space != .shared) return error.UnsupportedOperation;
917 const size = memref.size orelse return error.InvalidArtifact;
918 const bytes64 = std.math.mul(u64, size, @as(u64, elementByteSize(memref.element))) catch return error.InvalidArtifact;
919 const bytes = std.math.cast(u32, bytes64) orelse return error.InvalidArtifact;
920 if (bytes == 0) return error.InvalidArtifact;
921 const alignment64 = memref.alignment orelse elementByteSize(memref.element);
922 var alignment = std.math.cast(u32, alignment64) orelse return error.InvalidArtifact;
923 if (alignment == 0 or !std.math.isPowerOfTwo(alignment)) return error.InvalidArtifact;
924 if (memref.element == .f32) alignment = @max(alignment, 16);
925 const shared = self.allocShared();
926 if (op.getDynamicSize() != null) {
927 const byte_offset = try dynamicSharedByteOffset(op.op);
928 if (byte_offset % alignment != 0) return error.InvalidArtifact;
929 try self.shared_byte_offsets.put(self.allocator, shared, byte_offset);
930 self.dynamic_shared_alignment = @max(self.dynamic_shared_alignment, alignment);
931 try self.bind(result, .{ .shared = shared });
932 return;
933 }
934 try self.memory_decls.writer.print(" .shared .align {d} .b8 __choir_shared{d}[{d}];\n", .{ alignment, shared, bytes });
935 try self.bind(result, .{ .shared = shared });
936 }
937
938 fn emitAlloca(self: *Emitter, op: MemrefDialect.AllocaOp) EmitError!void {
939 if (op.getDynamicSize() != null) return error.UnsupportedOperation;
940 const result = op.getResult();
941 const info = memrefInfo(result.type) orelse return error.InvalidArtifact;
942 if (info.addr_space != .local) return error.UnsupportedOperation;
943 const size = info.size orelse return error.InvalidArtifact;
944 const bytes = std.math.mul(u64, size, elementByteSize(info.element)) catch
945 return error.InvalidArtifact;
946 if (bytes > std.math.maxInt(u32)) return error.UnsupportedOperation;
947 const alignment = info.alignment orelse elementByteSize(info.element);
948 if (alignment == 0 or !std.math.isPowerOfTwo(alignment)) return error.InvalidArtifact;
949 const ptr = self.allocPtr();
950 try self.memory_decls.writer.print(" .local .align {d} .b8 __choir_local{d}[{d}];\n", .{
951 alignment, ptr, @max(bytes, 1),
952 });
953 try self.line(" mov.u64 %rd{d}, __choir_local{d};", .{ ptr, ptr });
954 try self.bind(result, .{ .ptr = ptr });
955 }
956
957 fn emitNvptxLocalLoad(self: *Emitter, op: NvptxDialect.LoadLocalOp) EmitError!void {
958 const value = try self.require(op.getMemref());
959 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
960 if (info.addr_space != .local or value != .ptr) return error.InvalidArtifact;
961 try self.emitPointerLoad("local", op.getResult(), value.ptr, op.getIndex(), info.element);
962 }
963
964 fn emitNvptxLocalStore(self: *Emitter, op: NvptxDialect.StoreLocalOp) EmitError!void {
965 const value = try self.require(op.getMemref());
966 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
967 if (info.addr_space != .local or value != .ptr) return error.InvalidArtifact;
968 try self.emitPointerStore("local", op.getValue(), value.ptr, op.getIndex(), info.element);
969 }
970
971 fn emitPointerLoad(
972 self: *Emitter,
973 comptime space: []const u8,
974 result: *ir.Value,
975 ptr_reg: u32,
976 index: *ir.Value,
977 element: ScalarKind,
978 ) EmitError!void {
979 const addr = try self.emitPointerAddress(ptr_reg, index, element);
980 switch (element) {
981 .i8, .u8, .i16, .u16, .i32, .u32, .index => {
982 const out = self.allocU32();
983 const suffix = switch (element) {
984 .i8 => "s8",
985 .i16 => "s16",
986 .i32, .index => "u32",
987 else => @tagName(element),
988 };
989 try self.line(" ld." ++ space ++ ".{s} %r{d}, [%rd{d}+{d}];", .{
990 suffix, out, addr.reg, addr.imm,
991 });
992 try self.bind(result, intValueForKind(element, out));
993 },
994 .i64, .u64 => {
995 const out = self.allocU64();
996 try self.line(" ld." ++ space ++ ".u64 %rd{d}, [%rd{d}+{d}];", .{ out, addr.reg, addr.imm });
997 try self.bind(result, .{ .u64 = out });
998 },
999 .f32 => {
1000 const out = self.allocF32();
1001 try self.line(" ld." ++ space ++ ".f32 %f{d}, [%rd{d}+{d}];", .{ out, addr.reg, addr.imm });
1002 try self.bind(result, .{ .f32 = out });
1003 },
1004 .f64 => {
1005 const out = self.allocF64();
1006 try self.line(" ld." ++ space ++ ".f64 %fd{d}, [%rd{d}+{d}];", .{ out, addr.reg, addr.imm });
1007 try self.bind(result, .{ .f64 = out });
1008 },
1009 .f16 => {
1010 const half = self.allocB16();
1011 const out = self.allocF32();
1012 try self.line(" ld." ++ space ++ ".b16 %h{d}, [%rd{d}+{d}];", .{ half, addr.reg, addr.imm });
1013 try self.line(" cvt.f32.f16 %f{d}, %h{d};", .{ out, half });
1014 try self.bind(result, .{ .f32 = out });
1015 },
1016 .bf16 => {
1017 const bits = self.allocU32();
1018 try self.line(" ld." ++ space ++ ".b16 %r{d}, [%rd{d}+{d}];", .{ bits, addr.reg, addr.imm });
1019 try self.bind(result, .{ .f32 = try self.emitBf16BitsToF32(bits) });
1020 },
1021 .bool => {
1022 const byte = self.allocU32();
1023 const out = self.allocPred();
1024 try self.line(" ld." ++ space ++ ".u8 %r{d}, [%rd{d}+{d}];", .{ byte, addr.reg, addr.imm });
1025 try self.line(" setp.ne.u32 %p{d}, %r{d}, 0;", .{ out, byte });
1026 try self.bind(result, .{ .pred = out });
1027 },
1028 }
1029 }
1030
1031 fn emitSharedLoad(self: *Emitter, result: *ir.Value, shared: u32, index: *ir.Value, element: ScalarKind) EmitError!void {
1032 const addr = try self.emitSharedAddress(shared, index, element);
1033 switch (element) {
1034 .i8 => {
1035 const out = self.allocU32();
1036 try self.line(" ld.shared.s8 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1037 try self.bind(result, .{ .s32 = out });
1038 },
1039 .u8 => {
1040 const out = self.allocU32();
1041 try self.line(" ld.shared.u8 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1042 try self.bind(result, .{ .u32 = out });
1043 },
1044 .i16 => {
1045 const out = self.allocU32();
1046 try self.line(" ld.shared.s16 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1047 try self.bind(result, .{ .s32 = out });
1048 },
1049 .u16 => {
1050 const out = self.allocU32();
1051 try self.line(" ld.shared.u16 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1052 try self.bind(result, .{ .u32 = out });
1053 },
1054 .i64, .u64 => {
1055 const out = self.allocU64();
1056 try self.line(" ld.shared.u64 %rd{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1057 try self.bind(result, .{ .u64 = out });
1058 },
1059 .index, .u32 => {
1060 const out = self.allocU32();
1061 try self.line(" ld.shared.u32 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1062 try self.bind(result, .{ .u32 = out });
1063 },
1064 .i32 => {
1065 const out = self.allocU32();
1066 try self.line(" ld.shared.u32 %r{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1067 try self.bind(result, .{ .s32 = out });
1068 },
1069 .f32 => {
1070 const out = self.allocF32();
1071 try self.line(" ld.shared.f32 %f{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1072 try self.bind(result, .{ .f32 = out });
1073 },
1074 .f64 => {
1075 const out = self.allocF64();
1076 try self.line(" ld.shared.f64 %fd{d}, [%r{d}+{d}];", .{ out, addr.reg, addr.imm });
1077 try self.bind(result, .{ .f64 = out });
1078 },
1079 .f16 => {
1080 const half = self.allocB16();
1081 const out = self.allocF32();
1082 try self.line(" ld.shared.b16 %h{d}, [%r{d}+{d}];", .{ half, addr.reg, addr.imm });
1083 try self.line(" cvt.f32.f16 %f{d}, %h{d};", .{ out, half });
1084 try self.bind(result, .{ .f32 = out });
1085 },
1086 .bf16 => {
1087 const bits = self.allocU32();
1088 try self.line(" ld.shared.b16 %r{d}, [%r{d}+{d}];", .{ bits, addr.reg, addr.imm });
1089 try self.bind(result, .{ .f32 = try self.emitBf16BitsToF32(bits) });
1090 },
1091 .bool => {
1092 const byte = self.allocU32();
1093 const out = self.allocPred();
1094 try self.line(" ld.shared.u8 %r{d}, [%r{d}+{d}];", .{ byte, addr.reg, addr.imm });
1095 try self.line(" setp.ne.u32 %p{d}, %r{d}, 0;", .{ out, byte });
1096 try self.bind(result, .{ .pred = out });
1097 },
1098 }
1099 }
1100
1101 fn emitNvptxGlobalLoad(self: *Emitter, op: NvptxDialect.LoadGlobalOp) EmitError!void {
1102 const memref = try self.require(op.getMemref());
1103 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1104 switch (memref) {
1105 .ptr => |ptr_reg| {
1106 if (!isKernelParameterAddressSpace(info.addr_space)) return error.InvalidArtifact;
1107 if (isVec4F32Type(op.getResult().type)) {
1108 if (info.element != .f32) return error.UnsupportedOperation;
1109 const addr = try self.emitPointerAddress(ptr_reg, op.getIndex(), info.element);
1110 const base = self.allocF32x4();
1111 try self.line(" ld.global.v4.f32 {{%f{d}, %f{d}, %f{d}, %f{d}}}, [%rd{d}+{d}];", .{ base, base + 1, base + 2, base + 3, addr.reg, addr.imm });
1112 try self.bind(op.getResult(), .{ .f32x4 = base });
1113 return;
1114 }
1115 try self.emitPointerLoad("global", op.getResult(), ptr_reg, op.getIndex(), info.element);
1116 },
1117 else => return error.InvalidArtifact,
1118 }
1119 }
1120
1121 fn emitNvptxSharedLoad(self: *Emitter, op: NvptxDialect.LoadSharedOp) EmitError!void {
1122 const memref = try self.require(op.getMemref());
1123 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1124 switch (memref) {
1125 .shared => |shared| {
1126 if (info.addr_space != .shared) return error.InvalidArtifact;
1127 if (isVec4F32Type(op.getResult().type)) {
1128 if (info.element != .f32) return error.UnsupportedOperation;
1129 const addr = try self.emitSharedAddress(shared, op.getIndex(), info.element);
1130 const base = self.allocF32x4();
1131 try self.line(" ld.shared.v4.f32 {{%f{d}, %f{d}, %f{d}, %f{d}}}, [%r{d}+{d}];", .{ base, base + 1, base + 2, base + 3, addr.reg, addr.imm });
1132 try self.bind(op.getResult(), .{ .f32x4 = base });
1133 return;
1134 }
1135 try self.emitSharedLoad(op.getResult(), shared, op.getIndex(), info.element);
1136 },
1137 else => return error.InvalidArtifact,
1138 }
1139 }
1140
1141 fn emitPointerStore(
1142 self: *Emitter,
1143 comptime space: []const u8,
1144 value: *ir.Value,
1145 ptr_reg: u32,
1146 index: *ir.Value,
1147 element: ScalarKind,
1148 ) EmitError!void {
1149 const addr = try self.emitPointerAddress(ptr_reg, index, element);
1150 switch (element) {
1151 .i8, .u8 => try self.line(" st." ++ space ++ ".u8 [%rd{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1152 .i16, .u16 => try self.line(" st." ++ space ++ ".u16 [%rd{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1153 .i64, .u64 => try self.line(" st." ++ space ++ ".u64 [%rd{d}+{d}], %rd{d};", .{ addr.reg, addr.imm, try self.asU64(try self.require(value)) }),
1154 .i32, .u32, .index => try self.line(" st." ++ space ++ ".u32 [%rd{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1155 .f32 => try self.line(" st." ++ space ++ ".f32 [%rd{d}+{d}], %f{d};", .{ addr.reg, addr.imm, try self.asF32(try self.require(value)) }),
1156 .f64 => try self.line(" st." ++ space ++ ".f64 [%rd{d}+{d}], %fd{d};", .{ addr.reg, addr.imm, try self.asF64(try self.require(value)) }),
1157 .f16 => {
1158 const half = self.allocB16();
1159 try self.line(" cvt.rn.f16.f32 %h{d}, %f{d};", .{ half, try self.asF32(try self.require(value)) });
1160 try self.line(" st." ++ space ++ ".b16 [%rd{d}+{d}], %h{d};", .{ addr.reg, addr.imm, half });
1161 },
1162 .bf16 => try self.line(" st." ++ space ++ ".b16 [%rd{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.emitF32ToBf16Bits(try self.require(value)) }),
1163 .bool => try self.line(" st." ++ space ++ ".u8 [%rd{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.emitBoolByte(try self.require(value)) }),
1164 }
1165 }
1166
1167 fn emitSharedStore(self: *Emitter, value: *ir.Value, shared: u32, index: *ir.Value, element: ScalarKind) EmitError!void {
1168 const addr = try self.emitSharedAddress(shared, index, element);
1169 switch (element) {
1170 .i8, .u8 => try self.line(" st.shared.u8 [%r{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1171 .i16, .u16 => try self.line(" st.shared.u16 [%r{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1172 .i64, .u64 => try self.line(" st.shared.u64 [%r{d}+{d}], %rd{d};", .{ addr.reg, addr.imm, try self.asU64(try self.require(value)) }),
1173 .i32, .u32, .index => try self.line(" st.shared.u32 [%r{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.asU32(try self.require(value)) }),
1174 .f32 => try self.line(" st.shared.f32 [%r{d}+{d}], %f{d};", .{ addr.reg, addr.imm, try self.asF32(try self.require(value)) }),
1175 .f64 => try self.line(" st.shared.f64 [%r{d}+{d}], %fd{d};", .{ addr.reg, addr.imm, try self.asF64(try self.require(value)) }),
1176 .f16 => {
1177 const half = self.allocB16();
1178 try self.line(" cvt.rn.f16.f32 %h{d}, %f{d};", .{ half, try self.asF32(try self.require(value)) });
1179 try self.line(" st.shared.b16 [%r{d}+{d}], %h{d};", .{ addr.reg, addr.imm, half });
1180 },
1181 .bf16 => try self.line(" st.shared.b16 [%r{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.emitF32ToBf16Bits(try self.require(value)) }),
1182 .bool => try self.line(" st.shared.u8 [%r{d}+{d}], %r{d};", .{ addr.reg, addr.imm, try self.emitBoolByte(try self.require(value)) }),
1183 }
1184 }
1185
1186 fn emitNvptxGlobalStore(self: *Emitter, op: NvptxDialect.StoreGlobalOp) EmitError!void {
1187 const memref = try self.require(op.getMemref());
1188 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1189 switch (memref) {
1190 .ptr => |ptr_reg| {
1191 if (!isKernelParameterAddressSpace(info.addr_space)) return error.InvalidArtifact;
1192 const stored = try self.require(op.getValue());
1193 if (stored == .f32x4) {
1194 if (info.element != .f32) return error.UnsupportedOperation;
1195 const base = stored.f32x4;
1196 const addr = try self.emitPointerAddress(ptr_reg, op.getIndex(), info.element);
1197 try self.line(" st.global.v4.f32 [%rd{d}+{d}], {{%f{d}, %f{d}, %f{d}, %f{d}}};", .{ addr.reg, addr.imm, base, base + 1, base + 2, base + 3 });
1198 return;
1199 }
1200 try self.emitPointerStore("global", op.getValue(), ptr_reg, op.getIndex(), info.element);
1201 },
1202 else => return error.InvalidArtifact,
1203 }
1204 }
1205
1206 fn emitNvptxGlobalAtomic(self: *Emitter, op: NvptxDialect.AtomicGlobalOp) EmitError!void {
1207 const kind = op.getKind() orelse return error.InvalidArtifact;
1208 const memref = try self.require(op.getMemref());
1209 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1210 const ptr_reg = switch (memref) {
1211 .ptr => |ptr_reg| ptr_reg,
1212 else => return error.InvalidArtifact,
1213 };
1214 if (!isKernelParameterAddressSpace(info.addr_space)) return error.InvalidArtifact;
1215
1216 const addr = try self.emitPointerAddress(ptr_reg, op.getIndex(), info.element);
1217
1218 switch (info.element) {
1219 .f32 => {
1220 if (kind != .add) return error.UnsupportedOperation;
1221 const operand = try self.asF32(try self.require(op.getValue()));
1222 if (op.getResult().hasNoUses()) {
1223 try self.line(" red.global.add.f32 [%rd{d}+{d}], %f{d};", .{ addr.reg, addr.imm, operand });
1224 return;
1225 }
1226 const out = self.allocF32();
1227 try self.line(" atom.global.add.f32 %f{d}, [%rd{d}+{d}], %f{d};", .{ out, addr.reg, addr.imm, operand });
1228 try self.bind(op.getResult(), .{ .f32 = out });
1229 },
1230 .i32, .u32, .index => {
1231 const suffix = try atomicIntegerSuffix(kind, info.element);
1232 const operand = try self.asU32(try self.require(op.getValue()));
1233 if (op.getResult().hasNoUses() and atomicKindSupportsRed(kind)) {
1234 try self.line(" red.global.{s} [%rd{d}+{d}], %r{d};", .{ suffix, addr.reg, addr.imm, operand });
1235 return;
1236 }
1237 const out = self.allocU32();
1238 try self.line(" atom.global.{s} %r{d}, [%rd{d}+{d}], %r{d};", .{ suffix, out, addr.reg, addr.imm, operand });
1239 try self.bind(op.getResult(), intValueForKind(info.element, out));
1240 },
1241 .i8, .i16, .u8, .u16, .i64, .u64, .f16, .bf16, .f64, .bool => return error.UnsupportedOperation,
1242 }
1243 }
1244
1245 fn emitNvptxSharedAtomic(self: *Emitter, op: NvptxDialect.AtomicSharedOp) EmitError!void {
1246 const kind = op.getKind() orelse return error.InvalidArtifact;
1247 const memref = try self.require(op.getMemref());
1248 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1249 const shared = switch (memref) {
1250 .shared => |shared| shared,
1251 else => return error.InvalidArtifact,
1252 };
1253 if (info.addr_space != .shared) return error.InvalidArtifact;
1254
1255 const addr = try self.emitSharedAddress(shared, op.getIndex(), info.element);
1256
1257 switch (info.element) {
1258 .f32 => {
1259 if (kind != .add) return error.UnsupportedOperation;
1260 const operand = try self.asF32(try self.require(op.getValue()));
1261 if (op.getResult().hasNoUses()) {
1262 try self.line(" red.shared.add.f32 [%r{d}+{d}], %f{d};", .{ addr.reg, addr.imm, operand });
1263 return;
1264 }
1265 const out = self.allocF32();
1266 try self.line(" atom.shared.add.f32 %f{d}, [%r{d}+{d}], %f{d};", .{ out, addr.reg, addr.imm, operand });
1267 try self.bind(op.getResult(), .{ .f32 = out });
1268 },
1269 .i32, .u32, .index => {
1270 const suffix = try atomicIntegerSuffix(kind, info.element);
1271 const operand = try self.asU32(try self.require(op.getValue()));
1272 if (op.getResult().hasNoUses() and atomicKindSupportsRed(kind)) {
1273 try self.line(" red.shared.{s} [%r{d}+{d}], %r{d};", .{ suffix, addr.reg, addr.imm, operand });
1274 return;
1275 }
1276 const out = self.allocU32();
1277 try self.line(" atom.shared.{s} %r{d}, [%r{d}+{d}], %r{d};", .{ suffix, out, addr.reg, addr.imm, operand });
1278 try self.bind(op.getResult(), intValueForKind(info.element, out));
1279 },
1280 .i8, .i16, .u8, .u16, .i64, .u64, .f16, .bf16, .f64, .bool => return error.UnsupportedOperation,
1281 }
1282 }
1283
1284 fn emitNvptxGlobalAtomicCas(self: *Emitter, op: NvptxDialect.AtomicCasGlobalOp) EmitError!void {
1285 const memref = try self.require(op.getMemref());
1286 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1287 const ptr_reg = switch (memref) {
1288 .ptr => |ptr_reg| ptr_reg,
1289 else => return error.InvalidArtifact,
1290 };
1291 if (!isKernelParameterAddressSpace(info.addr_space)) return error.InvalidArtifact;
1292
1293 const addr = try self.emitPointerAddress(ptr_reg, op.getIndex(), info.element);
1294
1295 switch (info.element) {
1296 .i32, .u32, .index => {
1297 const expected = try self.asU32(try self.require(op.getExpected()));
1298 const desired = try self.asU32(try self.require(op.getDesired()));
1299 const out = self.allocU32();
1300 try self.line(" atom.global.cas.b32 %r{d}, [%rd{d}+{d}], %r{d}, %r{d};", .{ out, addr.reg, addr.imm, expected, desired });
1301 try self.bind(op.getResult(), intValueForKind(info.element, out));
1302 },
1303 .i8, .i16, .u8, .u16, .i64, .u64, .f16, .bf16, .f32, .f64, .bool => return error.UnsupportedOperation,
1304 }
1305 }
1306
1307 fn emitNvptxSharedAtomicCas(self: *Emitter, op: NvptxDialect.AtomicCasSharedOp) EmitError!void {
1308 const memref = try self.require(op.getMemref());
1309 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1310 const shared = switch (memref) {
1311 .shared => |shared| shared,
1312 else => return error.InvalidArtifact,
1313 };
1314 if (info.addr_space != .shared) return error.InvalidArtifact;
1315
1316 const addr = try self.emitSharedAddress(shared, op.getIndex(), info.element);
1317
1318 switch (info.element) {
1319 .i32, .u32, .index => {
1320 const expected = try self.asU32(try self.require(op.getExpected()));
1321 const desired = try self.asU32(try self.require(op.getDesired()));
1322 const out = self.allocU32();
1323 try self.line(" atom.shared.cas.b32 %r{d}, [%r{d}+{d}], %r{d}, %r{d};", .{ out, addr.reg, addr.imm, expected, desired });
1324 try self.bind(op.getResult(), intValueForKind(info.element, out));
1325 },
1326 .i8, .i16, .u8, .u16, .i64, .u64, .f16, .bf16, .f32, .f64, .bool => return error.UnsupportedOperation,
1327 }
1328 }
1329
1330 fn emitNvptxSharedStore(self: *Emitter, op: NvptxDialect.StoreSharedOp) EmitError!void {
1331 const memref = try self.require(op.getMemref());
1332 const info = memrefInfo(op.getMemref().type) orelse return error.InvalidArtifact;
1333 switch (memref) {
1334 .shared => |shared| {
1335 if (info.addr_space != .shared) return error.InvalidArtifact;
1336 const stored = try self.require(op.getValue());
1337 if (stored == .f32x4) {
1338 if (info.element != .f32) return error.UnsupportedOperation;
1339 const base = stored.f32x4;
1340 const addr = try self.emitSharedAddress(shared, op.getIndex(), info.element);
1341 try self.line(" st.shared.v4.f32 [%r{d}+{d}], {{%f{d}, %f{d}, %f{d}, %f{d}}};", .{ addr.reg, addr.imm, base, base + 1, base + 2, base + 3 });
1342 return;
1343 }
1344 try self.emitSharedStore(op.getValue(), shared, op.getIndex(), info.element);
1345 },
1346 else => return error.InvalidArtifact,
1347 }
1348 }
1349
1350 fn emitBoolByte(self: *Emitter, value: Value) EmitError!u32 {
1351 const pred = try self.asPred(value);
1352 const one = self.allocU32();
1353 const zero = self.allocU32();
1354 const out = self.allocU32();
1355 try self.line(" mov.u32 %r{d}, 1;", .{one});
1356 try self.line(" mov.u32 %r{d}, 0;", .{zero});
1357 try self.line(" selp.u32 %r{d}, %r{d}, %r{d}, %p{d};", .{ out, one, zero, pred });
1358 return out;
1359 }
1360
1361 fn emitSharedAddress(self: *Emitter, shared: u32, index: *ir.Value, element: ScalarKind) EmitError!AddressParts {
1362 const bytes: u32 = elementByteSize(element);
1363 const shared_byte_offset = self.shared_byte_offsets.get(shared) orelse 0;
1364 var peeled = peelIndexOffset(index);
1365 var imm = std.math.mul(i64, peeled.offset, @as(i64, bytes)) catch return error.InvalidArtifact;
1366 imm = std.math.add(i64, imm, @as(i64, shared_byte_offset)) catch return error.InvalidArtifact;
1367 if (imm > max_address_immediate or imm < -max_address_immediate) {
1368 peeled = .{ .base = index, .offset = 0 };
1369 imm = shared_byte_offset;
1370 }
1371 if (imm > max_address_immediate or imm < -max_address_immediate) return error.InvalidArtifact;
1372 const key = AddressBaseKey{
1373 .space = shared,
1374 .base = if (peeled.base) |base| @intFromPtr(base) else 0,
1375 .bytes = bytes,
1376 };
1377 if (self.shared_bases.get(key)) |reg| return .{ .reg = reg, .imm = imm };
1378 const addr = self.allocU32();
1379 if (peeled.base) |base_value| {
1380 const base_reg = try self.asU32(try self.require(base_value));
1381 const offset_reg = self.allocU32();
1382 try self.line(" mul.lo.u32 %r{d}, %r{d}, {d};", .{ offset_reg, base_reg, bytes });
1383 try self.emitSharedBaseMove(addr, shared);
1384 try self.line(" add.u32 %r{d}, %r{d}, %r{d};", .{ addr, addr, offset_reg });
1385 } else {
1386 try self.emitSharedBaseMove(addr, shared);
1387 }
1388 self.shared_bases.put(self.allocator, key, addr) catch return error.OutOfMemory;
1389 return .{ .reg = addr, .imm = imm };
1390 }
1391
1392 fn emitSharedBaseMove(self: *Emitter, addr: u32, shared: u32) EmitError!void {
1393 if (self.shared_byte_offsets.contains(shared)) {
1394 try self.line(" mov.u32 %r{d}, __choir_dynamic_shared;", .{addr});
1395 } else {
1396 try self.line(" mov.u32 %r{d}, __choir_shared{d};", .{ addr, shared });
1397 }
1398 }
1399
1400 fn emitPointerAddress(self: *Emitter, ptr_reg: u32, index: *ir.Value, element: ScalarKind) EmitError!AddressParts {
1401 const bytes: u32 = elementByteSize(element);
1402 var peeled = peelIndexOffset(index);
1403 var imm = std.math.mul(i64, peeled.offset, bytes) catch overflow: {
1404 peeled = .{ .base = index, .offset = 0 };
1405 break :overflow 0;
1406 };
1407 if (imm > max_address_immediate or imm < -max_address_immediate) {
1408 peeled = .{ .base = index, .offset = 0 };
1409 imm = 0;
1410 }
1411 const base_value = peeled.base orelse return .{ .reg = ptr_reg, .imm = imm };
1412 const key = AddressBaseKey{
1413 .space = ptr_reg,
1414 .base = @intFromPtr(base_value),
1415 .bytes = bytes,
1416 };
1417 if (self.pointer_bases.get(key)) |reg| return .{ .reg = reg, .imm = imm };
1418 const base = try self.require(base_value);
1419 const offset_reg = self.allocPtr();
1420 const addr = self.allocPtr();
1421 switch (base) {
1422 .u64 => |reg| try self.line(" mul.lo.u64 %rd{d}, %rd{d}, {d};", .{
1423 offset_reg, reg, bytes,
1424 }),
1425 .s32 => |reg| try self.line(" mul.wide.s32 %rd{d}, %r{d}, {d};", .{
1426 offset_reg, reg, bytes,
1427 }),
1428 .u32 => |reg| try self.line(" mul.wide.u32 %rd{d}, %r{d}, {d};", .{
1429 offset_reg, reg, bytes,
1430 }),
1431 else => return error.UnsupportedOperation,
1432 }
1433 try self.line(" add.s64 %rd{d}, %rd{d}, %rd{d};", .{ addr, ptr_reg, offset_reg });
1434 self.pointer_bases.put(self.allocator, key, addr) catch return error.OutOfMemory;
1435 return .{ .reg = addr, .imm = imm };
1436 }
1437
1438 fn emitConstant(self: *Emitter, op: ArithDialect.ConstantOp) EmitError!void {
1439 const result = op.getResult();
1440 switch (try computeKind(result.type)) {
1441 .index => {
1442 const value = op.getIntValue() orelse return error.InvalidArtifact;
1443 const out = self.allocU32();
1444 try self.line(" mov.u32 %r{d}, {d};", .{ out, @as(u32, @bitCast(@as(i32, @intCast(value)))) });
1445 try self.bind(result, .{ .u32 = out });
1446 },
1447 .i8, .i16, .i32 => {
1448 const value = op.getIntValue() orelse return error.InvalidArtifact;
1449 const out = self.allocU32();
1450 try self.line(" mov.u32 %r{d}, {d};", .{ out, @as(u32, @bitCast(@as(i32, @intCast(value)))) });
1451 try self.bind(result, .{ .s32 = out });
1452 },
1453 .u8, .u16, .u32 => {
1454 const value = op.getIntValue() orelse return error.InvalidArtifact;
1455 const out = self.allocU32();
1456 const narrowed = std.math.cast(u32, value) orelse return error.InvalidArtifact;
1457 try self.line(" mov.u32 %r{d}, {d};", .{ out, narrowed });
1458 try self.bind(result, .{ .u32 = out });
1459 },
1460 .i64 => {
1461 const value = op.getIntValue() orelse return error.InvalidArtifact;
1462 const out = self.allocU64();
1463 try self.line(" mov.u64 %rd{d}, {d};", .{ out, @as(u64, @bitCast(value)) });
1464 try self.bind(result, .{ .u64 = out });
1465 },
1466 .u64 => {
1467 const value = op.getIntValue() orelse return error.InvalidArtifact;
1468 const out = self.allocU64();
1469 const narrowed = std.math.cast(u64, value) orelse return error.InvalidArtifact;
1470 try self.line(" mov.u64 %rd{d}, {d};", .{ out, narrowed });
1471 try self.bind(result, .{ .u64 = out });
1472 },
1473 .f32 => {
1474 const value = op.getFloatValue() orelse return error.InvalidArtifact;
1475 const out = self.allocF32();
1476 try self.line(" mov.f32 %f{d}, 0f{X:0>8};", .{ out, f32Bits(value) });
1477 try self.bind(result, .{ .f32 = out });
1478 },
1479 .f64 => {
1480 const value = op.getFloatValue() orelse return error.InvalidArtifact;
1481 const out = self.allocF64();
1482 try self.line(" mov.f64 %fd{d}, 0d{X:0>16};", .{ out, f64Bits(value) });
1483 try self.bind(result, .{ .f64 = out });
1484 },
1485 .f16, .bf16 => unreachable,
1486 .bool => {
1487 const bool_attr = op.op.getAttrAs(ir.Attribute.BoolAttr, "value") orelse return error.InvalidArtifact;
1488 const out = self.allocPred();
1489 const source: u32 = if (bool_attr.getValue()) 1 else 0;
1490 try self.line(" setp.ne.u32 %p{d}, {d}, 0;", .{ out, source });
1491 try self.bind(result, .{ .pred = out });
1492 },
1493 }
1494 }
1495
1496 fn emitBinary(self: *Emitter, op: *ir.Operation, comptime instruction: []const u8) EmitError!void {
1497 const result = op.getResult(0) orelse return error.InvalidArtifact;
1498 const lhs = try self.require(op.operands.items[0].value);
1499 const rhs = try self.require(op.operands.items[1].value);
1500 switch (try computeKind(result.type)) {
1501 .index, .u8, .u16, .u32 => {
1502 const out = self.allocU32();
1503 if (std.mem.eql(u8, instruction, "mul")) {
1504 try self.line(" mul.lo.u32 %r{d}, %r{d}, %r{d};", .{ out, try self.asU32(lhs), try self.asU32(rhs) });
1505 } else if (std.mem.eql(u8, instruction, "max") or std.mem.eql(u8, instruction, "min")) {
1506 try self.line(" {s}.u32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1507 } else {
1508 try self.line(" {s}.u32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1509 }
1510 try self.bind(result, .{ .u32 = out });
1511 },
1512 .i8, .i16, .i32 => {
1513 const out = self.allocU32();
1514 if (std.mem.eql(u8, instruction, "mul")) {
1515 try self.line(" mul.lo.u32 %r{d}, %r{d}, %r{d};", .{ out, try self.asU32(lhs), try self.asU32(rhs) });
1516 } else if (std.mem.eql(u8, instruction, "max") or std.mem.eql(u8, instruction, "min")) {
1517 try self.line(" {s}.s32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1518 } else {
1519 try self.line(" {s}.u32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1520 }
1521 try self.bind(result, .{ .s32 = out });
1522 },
1523 .i64 => {
1524 const out = self.allocU64();
1525 if (std.mem.eql(u8, instruction, "mul")) {
1526 try self.line(" mul.lo.u64 %rd{d}, %rd{d}, %rd{d};", .{ out, try self.asU64(lhs), try self.asU64(rhs) });
1527 } else if (std.mem.eql(u8, instruction, "max") or std.mem.eql(u8, instruction, "min")) {
1528 try self.line(" {s}.s64 %rd{d}, %rd{d}, %rd{d};", .{ instruction, out, try self.asU64(lhs), try self.asU64(rhs) });
1529 } else {
1530 try self.line(" {s}.u64 %rd{d}, %rd{d}, %rd{d};", .{ instruction, out, try self.asU64(lhs), try self.asU64(rhs) });
1531 }
1532 try self.bind(result, .{ .u64 = out });
1533 },
1534 .u64 => {
1535 const out = self.allocU64();
1536 if (std.mem.eql(u8, instruction, "mul")) {
1537 try self.line(" mul.lo.u64 %rd{d}, %rd{d}, %rd{d};", .{ out, try self.asU64(lhs), try self.asU64(rhs) });
1538 } else if (std.mem.eql(u8, instruction, "max") or std.mem.eql(u8, instruction, "min")) {
1539 try self.line(" {s}.u64 %rd{d}, %rd{d}, %rd{d};", .{ instruction, out, try self.asU64(lhs), try self.asU64(rhs) });
1540 } else {
1541 try self.line(" {s}.u64 %rd{d}, %rd{d}, %rd{d};", .{ instruction, out, try self.asU64(lhs), try self.asU64(rhs) });
1542 }
1543 try self.bind(result, .{ .u64 = out });
1544 },
1545 .f32 => {
1546 const out = self.allocF32();
1547 try self.line(" {s}.f32 %f{d}, %f{d}, %f{d};", .{ instruction, out, try self.asF32(lhs), try self.asF32(rhs) });
1548 try self.bind(result, .{ .f32 = out });
1549 },
1550 .f64 => {
1551 const out = self.allocF64();
1552 try self.line(" {s}.f64 %fd{d}, %fd{d}, %fd{d};", .{ instruction, out, try self.asF64(lhs), try self.asF64(rhs) });
1553 try self.bind(result, .{ .f64 = out });
1554 },
1555 .f16, .bf16 => unreachable,
1556 .bool => return error.UnsupportedOperation,
1557 }
1558 }
1559
1560 fn emitUmulhi(self: *Emitter, op: *ir.Operation) EmitError!void {
1561 const result = op.getResult(0) orelse return error.InvalidArtifact;
1562 const lhs = try self.require(op.operands.items[0].value);
1563 const rhs = try self.require(op.operands.items[1].value);
1564 switch (try computeKind(result.type)) {
1565 .index, .i8, .i16, .i32, .u8, .u16, .u32 => {
1566 const out = self.allocU32();
1567 try self.line(" mul.hi.u32 %r{d}, %r{d}, %r{d};", .{ out, try self.asU32(lhs), try self.asU32(rhs) });
1568 try self.bind(result, intValueForKind(try computeKind(result.type), out));
1569 },
1570 .i64, .u64 => {
1571 const out = self.allocU64();
1572 try self.line(" mul.hi.u64 %rd{d}, %rd{d}, %rd{d};", .{ out, try self.asU64(lhs), try self.asU64(rhs) });
1573 try self.bind(result, .{ .u64 = out });
1574 },
1575 .f32, .f16, .bf16, .f64, .bool => return error.UnsupportedOperation,
1576 }
1577 }
1578
1579 fn emitNeg(self: *Emitter, op: *ir.Operation) EmitError!void {
1580 const result = op.getResult(0) orelse return error.InvalidArtifact;
1581 const input = try self.require(op.operands.items[0].value);
1582 switch (try computeKind(result.type)) {
1583 .i64, .u64 => {
1584 const out = self.allocU64();
1585 try self.line(" neg.s64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) });
1586 try self.bind(result, .{ .u64 = out });
1587 },
1588 .i8, .i16, .i32 => {
1589 const out = self.allocU32();
1590 try self.line(" neg.s32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
1591 try self.bind(result, .{ .s32 = out });
1592 },
1593 .u8, .u16, .u32, .index => {
1594 const out = self.allocU32();
1595 try self.line(" neg.s32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
1596 try self.bind(result, .{ .u32 = out });
1597 },
1598 .f32 => {
1599 const out = self.allocF32();
1600 try self.line(" neg.f32 %f{d}, %f{d};", .{ out, try self.asF32(input) });
1601 try self.bind(result, .{ .f32 = out });
1602 },
1603 .f64 => {
1604 const out = self.allocF64();
1605 try self.line(" neg.f64 %fd{d}, %fd{d};", .{ out, try self.asF64(input) });
1606 try self.bind(result, .{ .f64 = out });
1607 },
1608 .f16, .bf16 => unreachable,
1609 .bool => return error.UnsupportedOperation,
1610 }
1611 }
1612
1613 fn emitDiv(self: *Emitter, op: *ir.Operation) EmitError!void {
1614 const result = op.getResult(0) orelse return error.InvalidArtifact;
1615 const lhs = try self.require(op.operands.items[0].value);
1616 const rhs = try self.require(op.operands.items[1].value);
1617 switch (try computeKind(result.type)) {
1618 .i64 => {
1619 const out = self.allocU64();
1620 try self.line(" div.s64 %rd{d}, %rd{d}, %rd{d};", .{ out, try self.asU64(lhs), try self.asU64(rhs) });
1621 try self.bind(result, .{ .u64 = out });
1622 },
1623 .u64 => {
1624 const out = self.allocU64();
1625 try self.line(" div.u64 %rd{d}, %rd{d}, %rd{d};", .{ out, try self.asU64(lhs), try self.asU64(rhs) });
1626 try self.bind(result, .{ .u64 = out });
1627 },
1628 .index, .u8, .u16, .u32 => {
1629 const out = self.allocU32();
1630 try self.line(" div.u32 %r{d}, %r{d}, %r{d};", .{ out, try self.asU32(lhs), try self.asU32(rhs) });
1631 try self.bind(result, .{ .u32 = out });
1632 },
1633 .i8, .i16, .i32 => {
1634 const out = self.allocU32();
1635 try self.line(" div.s32 %r{d}, %r{d}, %r{d};", .{ out, try self.asU32(lhs), try self.asU32(rhs) });
1636 try self.bind(result, .{ .s32 = out });
1637 },
1638 .f32 => {
1639 const out = self.allocF32();
1640 try self.line(" div.approx.f32 %f{d}, %f{d}, %f{d};", .{ out, try self.asF32(lhs), try self.asF32(rhs) });
1641 try self.bind(result, .{ .f32 = out });
1642 },
1643 .f64 => {
1644 const out = self.allocF64();
1645 try self.line(" div.rn.f64 %fd{d}, %fd{d}, %fd{d};", .{ out, try self.asF64(lhs), try self.asF64(rhs) });
1646 try self.bind(result, .{ .f64 = out });
1647 },
1648 .f16, .bf16 => unreachable,
1649 .bool => return error.UnsupportedOperation,
1650 }
1651 }
1652
1653 fn emitAbs(self: *Emitter, op: *ir.Operation) EmitError!void {
1654 const result = op.getResult(0) orelse return error.InvalidArtifact;
1655 const input = try self.require(op.operands.items[0].value);
1656 switch (try computeKind(result.type)) {
1657 .i64 => {
1658 const out = self.allocU64();
1659 try self.line(" abs.s64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) });
1660 try self.bind(result, .{ .u64 = out });
1661 },
1662 .u64 => {
1663 const out = self.allocU64();
1664 try self.line(" mov.u64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) });
1665 try self.bind(result, .{ .u64 = out });
1666 },
1667 .i8, .i16, .i32 => {
1668 const out = self.allocU32();
1669 try self.line(" abs.s32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
1670 try self.bind(result, .{ .s32 = out });
1671 },
1672 .u8, .u16, .u32, .index => {
1673 const out = self.allocU32();
1674 try self.line(" mov.u32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
1675 try self.bind(result, .{ .u32 = out });
1676 },
1677 .f32 => {
1678 const out = self.allocF32();
1679 try self.line(" abs.f32 %f{d}, %f{d};", .{ out, try self.asF32(input) });
1680 try self.bind(result, .{ .f32 = out });
1681 },
1682 .f64 => {
1683 const out = self.allocF64();
1684 try self.line(" abs.f64 %fd{d}, %fd{d};", .{ out, try self.asF64(input) });
1685 try self.bind(result, .{ .f64 = out });
1686 },
1687 .f16, .bf16 => unreachable,
1688 .bool => return error.UnsupportedOperation,
1689 }
1690 }
1691
1692 fn emitFloatUnary(self: *Emitter, op: *ir.Operation, comptime instruction: []const u8) EmitError!void {
1693 const result = op.getResult(0) orelse return error.InvalidArtifact;
1694 if ((try computeKind(result.type)) != .f32) return error.UnsupportedOperation;
1695 const input = try self.asF32(try self.require(op.operands.items[0].value));
1696 const out = self.allocF32();
1697 try self.line(" {s}.f32 %f{d}, %f{d};", .{ instruction, out, input });
1698 try self.bind(result, .{ .f32 = out });
1699 }
1700
1701 fn emitFloatIntegralUnary(self: *Emitter, op: *ir.Operation, comptime instruction: []const u8) EmitError!void {
1702 const result = op.getResult(0) orelse return error.InvalidArtifact;
1703 if ((try computeKind(result.type)) != .f32) return error.UnsupportedOperation;
1704 const input = try self.asF32(try self.require(op.operands.items[0].value));
1705 const out = self.allocF32();
1706 try self.line(" {s}.f32.f32 %f{d}, %f{d};", .{ instruction, out, input });
1707 try self.bind(result, .{ .f32 = out });
1708 }
1709
1710 fn emitRound(self: *Emitter, op: *ir.Operation) EmitError!void {
1711 const result = op.getResult(0) orelse return error.InvalidArtifact;
1712 if ((try computeKind(result.type)) != .f32) return error.UnsupportedOperation;
1713 const input = try self.asF32(try self.require(op.operands.items[0].value));
1714 const abs_value = self.allocF32();
1715 const shifted = self.allocF32();
1716 const rounded_abs = self.allocF32();
1717 const input_bits = self.allocU32();
1718 const sign_mask = self.allocU32();
1719 const sign_bits = self.allocU32();
1720 const rounded_bits = self.allocU32();
1721 const magnitude_mask = self.allocU32();
1722 const magnitude_bits = self.allocU32();
1723 const result_bits = self.allocU32();
1724 const out = self.allocF32();
1725 try self.line(" abs.f32 %f{d}, %f{d};", .{ abs_value, input });
1726 try self.line(" add.f32 %f{d}, %f{d}, 0f3F000000;", .{ shifted, abs_value });
1727 try self.line(" cvt.rmi.f32.f32 %f{d}, %f{d};", .{ rounded_abs, shifted });
1728 try self.line(" mov.b32 %r{d}, %f{d};", .{ input_bits, input });
1729 try self.line(" mov.u32 %r{d}, 2147483648;", .{sign_mask});
1730 try self.line(" and.b32 %r{d}, %r{d}, %r{d};", .{ sign_bits, input_bits, sign_mask });
1731 try self.line(" mov.b32 %r{d}, %f{d};", .{ rounded_bits, rounded_abs });
1732 try self.line(" mov.u32 %r{d}, 2147483647;", .{magnitude_mask});
1733 try self.line(" and.b32 %r{d}, %r{d}, %r{d};", .{ magnitude_bits, rounded_bits, magnitude_mask });
1734 try self.line(" or.b32 %r{d}, %r{d}, %r{d};", .{ result_bits, magnitude_bits, sign_bits });
1735 try self.line(" mov.b32 %f{d}, %r{d};", .{ out, result_bits });
1736 try self.bind(result, .{ .f32 = out });
1737 }
1738
1739 fn emitCompare(self: *Emitter, op: ArithDialect.CmpOp) EmitError!void {
1740 const result = op.getResult();
1741 const lhs = try self.require(op.op.operands.items[0].value);
1742 const rhs = try self.require(op.op.operands.items[1].value);
1743 const pred = op.getPredicate() orelse return error.InvalidArtifact;
1744 const out = self.allocPred();
1745 switch (try computeKind(op.op.operands.items[0].value.type)) {
1746 .index, .u8, .u16, .u32 => try self.line(" setp.{s}.u32 %p{d}, %r{d}, %r{d};", .{ ptxPredicate(pred), out, try self.asU32(lhs), try self.asU32(rhs) }),
1747 .i8, .i16, .i32 => try self.line(" setp.{s}.s32 %p{d}, %r{d}, %r{d};", .{ ptxPredicate(pred), out, try self.asU32(lhs), try self.asU32(rhs) }),
1748 .i64 => try self.line(" setp.{s}.s64 %p{d}, %rd{d}, %rd{d};", .{ ptxPredicate(pred), out, try self.asU64(lhs), try self.asU64(rhs) }),
1749 .u64 => try self.line(" setp.{s}.u64 %p{d}, %rd{d}, %rd{d};", .{ ptxPredicate(pred), out, try self.asU64(lhs), try self.asU64(rhs) }),
1750 .f32 => try self.line(" setp.{s}.f32 %p{d}, %f{d}, %f{d};", .{ ptxPredicate(pred), out, try self.asF32(lhs), try self.asF32(rhs) }),
1751 .f64 => try self.line(" setp.{s}.f64 %p{d}, %fd{d}, %fd{d};", .{ ptxPredicate(pred), out, try self.asF64(lhs), try self.asF64(rhs) }),
1752 .f16, .bf16 => unreachable,
1753 .bool => return error.UnsupportedOperation,
1754 }
1755 try self.bind(result, .{ .pred = out });
1756 }
1757
1758 fn emitSelect(self: *Emitter, op: *ir.Operation) EmitError!void {
1759 const result = op.getResult(0) orelse return error.InvalidArtifact;
1760 const cond = try self.asPred(try self.require(op.operands.items[0].value));
1761 const true_value = try self.require(op.operands.items[1].value);
1762 const false_value = try self.require(op.operands.items[2].value);
1763 switch (try computeKind(result.type)) {
1764 .index, .u8, .u16, .u32 => {
1765 const out = self.allocU32();
1766 try self.line(" selp.u32 %r{d}, %r{d}, %r{d}, %p{d};", .{ out, try self.asU32(true_value), try self.asU32(false_value), cond });
1767 try self.bind(result, .{ .u32 = out });
1768 },
1769 .i8, .i16, .i32 => {
1770 const out = self.allocU32();
1771 try self.line(" selp.u32 %r{d}, %r{d}, %r{d}, %p{d};", .{ out, try self.asU32(true_value), try self.asU32(false_value), cond });
1772 try self.bind(result, .{ .s32 = out });
1773 },
1774 .i64, .u64 => {
1775 const out = self.allocU64();
1776 try self.line(" selp.b64 %rd{d}, %rd{d}, %rd{d}, %p{d};", .{ out, try self.asU64(true_value), try self.asU64(false_value), cond });
1777 try self.bind(result, .{ .u64 = out });
1778 },
1779 .f32 => {
1780 const out = self.allocF32();
1781 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ out, try self.asF32(true_value), try self.asF32(false_value), cond });
1782 try self.bind(result, .{ .f32 = out });
1783 },
1784 .f64 => {
1785 const out = self.allocF64();
1786 try self.line(" selp.f64 %fd{d}, %fd{d}, %fd{d}, %p{d};", .{ out, try self.asF64(true_value), try self.asF64(false_value), cond });
1787 try self.bind(result, .{ .f64 = out });
1788 },
1789 .f16, .bf16 => unreachable,
1790 .bool => {
1791 const inverted = self.allocPred();
1792 const true_taken = self.allocPred();
1793 const false_taken = self.allocPred();
1794 const out = self.allocPred();
1795 try self.line(" not.pred %p{d}, %p{d};", .{ inverted, cond });
1796 try self.line(" and.pred %p{d}, %p{d}, %p{d};", .{ true_taken, try self.asPred(true_value), cond });
1797 try self.line(" and.pred %p{d}, %p{d}, %p{d};", .{ false_taken, try self.asPred(false_value), inverted });
1798 try self.line(" or.pred %p{d}, %p{d}, %p{d};", .{ out, true_taken, false_taken });
1799 try self.bind(result, .{ .pred = out });
1800 },
1801 }
1802 }
1803
1804 fn emitCast(self: *Emitter, op: ArithDialect.CastOp) EmitError!void {
1805 const input = try self.require(op.getInput());
1806 const result = op.getResult();
1807 switch (try computeKind(result.type)) {
1808 .index, .u8, .u16, .u32 => {
1809 const out = self.allocU32();
1810 switch (input) {
1811 .u32, .s32 => try self.line(" mov.u32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1812 .u64 => try self.line(" cvt.u32.u64 %r{d}, %rd{d};", .{ out, try self.asU64(input) }),
1813 .pred => |pred| {
1814 const one = self.allocU32();
1815 const zero = self.allocU32();
1816 try self.line(" mov.u32 %r{d}, 1;", .{one});
1817 try self.line(" mov.u32 %r{d}, 0;", .{zero});
1818 try self.line(" selp.u32 %r{d}, %r{d}, %r{d}, %p{d};", .{ out, one, zero, pred });
1819 },
1820 .f32 => try self.line(" cvt.rzi.u32.f32 %r{d}, %f{d};", .{ out, try self.asF32(input) }),
1821 .f64 => try self.line(" cvt.rzi.u32.f64 %r{d}, %fd{d};", .{ out, try self.asF64(input) }),
1822 else => return error.UnsupportedOperation,
1823 }
1824 try self.bind(result, .{ .u32 = out });
1825 },
1826 .i8, .i16, .i32 => {
1827 const out = self.allocU32();
1828 switch (input) {
1829 .u32, .s32 => try self.line(" mov.u32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1830 .u64 => try self.line(" cvt.u32.u64 %r{d}, %rd{d};", .{ out, try self.asU64(input) }),
1831 .pred => |pred| {
1832 const one = self.allocU32();
1833 const zero = self.allocU32();
1834 try self.line(" mov.u32 %r{d}, 1;", .{one});
1835 try self.line(" mov.u32 %r{d}, 0;", .{zero});
1836 try self.line(" selp.u32 %r{d}, %r{d}, %r{d}, %p{d};", .{ out, one, zero, pred });
1837 },
1838 .f32 => try self.line(" cvt.rzi.s32.f32 %r{d}, %f{d};", .{ out, try self.asF32(input) }),
1839 .f64 => try self.line(" cvt.rzi.s32.f64 %r{d}, %fd{d};", .{ out, try self.asF64(input) }),
1840 else => return error.UnsupportedOperation,
1841 }
1842 try self.bind(result, .{ .s32 = out });
1843 },
1844 .i64 => {
1845 const out = self.allocU64();
1846 switch (input) {
1847 .s32 => try self.line(" cvt.s64.s32 %rd{d}, %r{d};", .{ out, try self.asU32(input) }),
1848 .u32 => try self.line(" cvt.u64.u32 %rd{d}, %r{d};", .{ out, try self.asU32(input) }),
1849 .u64 => try self.line(" mov.u64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) }),
1850 else => return error.UnsupportedOperation,
1851 }
1852 try self.bind(result, .{ .u64 = out });
1853 },
1854 .u64 => {
1855 const out = self.allocU64();
1856 switch (input) {
1857 .s32 => try self.line(" cvt.u64.s32 %rd{d}, %r{d};", .{ out, try self.asU32(input) }),
1858 .u32 => try self.line(" cvt.u64.u32 %rd{d}, %r{d};", .{ out, try self.asU32(input) }),
1859 .u64 => try self.line(" mov.u64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) }),
1860 .pred => |pred| {
1861 const one = self.allocU64();
1862 const zero = self.allocU64();
1863 try self.line(" mov.u64 %rd{d}, 1;", .{one});
1864 try self.line(" mov.u64 %rd{d}, 0;", .{zero});
1865 try self.line(" selp.b64 %rd{d}, %rd{d}, %rd{d}, %p{d};", .{ out, one, zero, pred });
1866 },
1867 .f32 => try self.line(" cvt.rzi.u64.f32 %rd{d}, %f{d};", .{ out, try self.asF32(input) }),
1868 .f64 => try self.line(" cvt.rzi.u64.f64 %rd{d}, %fd{d};", .{ out, try self.asF64(input) }),
1869 else => return error.UnsupportedOperation,
1870 }
1871 try self.bind(result, .{ .u64 = out });
1872 },
1873 .f32 => {
1874 const out = self.allocF32();
1875 switch (input) {
1876 .u32 => try self.line(" cvt.rn.f32.u32 %f{d}, %r{d};", .{ out, try self.asU32(input) }),
1877 .s32 => try self.line(" cvt.rn.f32.s32 %f{d}, %r{d};", .{ out, try self.asU32(input) }),
1878 .u64 => try self.line(" cvt.rn.f32.u64 %f{d}, %rd{d};", .{ out, try self.asU64(input) }),
1879 .f32 => try self.line(" mov.f32 %f{d}, %f{d};", .{ out, try self.asF32(input) }),
1880 .f64 => try self.line(" cvt.rn.f32.f64 %f{d}, %fd{d};", .{ out, try self.asF64(input) }),
1881 else => return error.UnsupportedOperation,
1882 }
1883 try self.bind(result, .{ .f32 = out });
1884 },
1885 .f64 => {
1886 const out = self.allocF64();
1887 switch (input) {
1888 .u32 => try self.line(" cvt.rn.f64.u32 %fd{d}, %r{d};", .{ out, try self.asU32(input) }),
1889 .s32 => try self.line(" cvt.rn.f64.s32 %fd{d}, %r{d};", .{ out, try self.asU32(input) }),
1890 .u64 => try self.line(" cvt.rn.f64.u64 %fd{d}, %rd{d};", .{ out, try self.asU64(input) }),
1891 .f32 => try self.line(" cvt.rn.f64.f32 %fd{d}, %f{d};", .{ out, try self.asF32(input) }),
1892 .f64 => try self.line(" mov.f64 %fd{d}, %fd{d};", .{ out, try self.asF64(input) }),
1893 else => return error.UnsupportedOperation,
1894 }
1895 try self.bind(result, .{ .f64 = out });
1896 },
1897 .f16, .bf16 => unreachable,
1898 .bool => return error.UnsupportedOperation,
1899 }
1900 }
1901
1902 fn emitBitcast(self: *Emitter, op: *ir.Operation) EmitError!void {
1903 const result = op.getResult(0) orelse return error.InvalidArtifact;
1904 const input = try self.require(op.operands.items[0].value);
1905 switch (try scalarKind(result.type)) {
1906 .i8, .i16 => {
1907 const out = self.allocU32();
1908 switch (input) {
1909 .u32, .s32 => try self.line(" mov.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1910 else => return error.UnsupportedOperation,
1911 }
1912 try self.bind(result, .{ .s32 = out });
1913 },
1914 .u8, .u16 => {
1915 const out = self.allocU32();
1916 switch (input) {
1917 .u32, .s32 => try self.line(" mov.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1918 else => return error.UnsupportedOperation,
1919 }
1920 try self.bind(result, .{ .u32 = out });
1921 },
1922 .i64, .u64 => {
1923 const out = self.allocU64();
1924 switch (input) {
1925 .u64 => try self.line(" mov.b64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) }),
1926 .f64 => try self.line(" mov.b64 %rd{d}, %fd{d};", .{ out, try self.asF64(input) }),
1927 else => return error.UnsupportedOperation,
1928 }
1929 try self.bind(result, .{ .u64 = out });
1930 },
1931 .index, .u32 => {
1932 const out = self.allocU32();
1933 switch (input) {
1934 .u32, .s32 => try self.line(" mov.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1935 .f32 => try self.line(" mov.b32 %r{d}, %f{d};", .{ out, try self.asF32(input) }),
1936 else => return error.UnsupportedOperation,
1937 }
1938 try self.bind(result, .{ .u32 = out });
1939 },
1940 .i32 => {
1941 const out = self.allocU32();
1942 switch (input) {
1943 .u32, .s32 => try self.line(" mov.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) }),
1944 .f32 => try self.line(" mov.b32 %r{d}, %f{d};", .{ out, try self.asF32(input) }),
1945 else => return error.UnsupportedOperation,
1946 }
1947 try self.bind(result, .{ .s32 = out });
1948 },
1949 .f32 => {
1950 const out = self.allocF32();
1951 switch (input) {
1952 .u32, .s32 => try self.line(" mov.b32 %f{d}, %r{d};", .{ out, try self.asU32(input) }),
1953 .f32 => try self.line(" mov.b32 %f{d}, %f{d};", .{ out, try self.asF32(input) }),
1954 else => return error.UnsupportedOperation,
1955 }
1956 try self.bind(result, .{ .f32 = out });
1957 },
1958 .f64 => {
1959 const out = self.allocF64();
1960 switch (input) {
1961 .u64 => try self.line(" mov.b64 %fd{d}, %rd{d};", .{ out, try self.asU64(input) }),
1962 .f64 => try self.line(" mov.b64 %fd{d}, %fd{d};", .{ out, try self.asF64(input) }),
1963 else => return error.UnsupportedOperation,
1964 }
1965 try self.bind(result, .{ .f64 = out });
1966 },
1967 .f16, .bf16 => return error.UnsupportedOperation,
1968 .bool => return error.UnsupportedOperation,
1969 }
1970 }
1971
1972 fn emitBitwise(self: *Emitter, op: *ir.Operation, comptime instruction: []const u8) EmitError!void {
1973 const result = op.getResult(0) orelse return error.InvalidArtifact;
1974 const lhs = try self.require(op.operands.items[0].value);
1975 const rhs = try self.require(op.operands.items[1].value);
1976 switch (try scalarKind(result.type)) {
1977 .index, .u8, .u16, .u32 => {
1978 const out = self.allocU32();
1979 try self.line(" {s}.b32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1980 try self.bind(result, .{ .u32 = out });
1981 },
1982 .i8, .i16, .i32 => {
1983 const out = self.allocU32();
1984 try self.line(" {s}.b32 %r{d}, %r{d}, %r{d};", .{ instruction, out, try self.asU32(lhs), try self.asU32(rhs) });
1985 try self.bind(result, .{ .s32 = out });
1986 },
1987 .i64, .u64 => {
1988 const out = self.allocU64();
1989 try self.line(" {s}.b64 %rd{d}, %rd{d}, %rd{d};", .{ instruction, out, try self.asU64(lhs), try self.asU64(rhs) });
1990 try self.bind(result, .{ .u64 = out });
1991 },
1992 .f16, .bf16 => return error.UnsupportedOperation,
1993 .bool => {
1994 const out = self.allocPred();
1995 try self.line(" {s}.pred %p{d}, %p{d}, %p{d};", .{ instruction, out, try self.asPred(lhs), try self.asPred(rhs) });
1996 try self.bind(result, .{ .pred = out });
1997 },
1998 .f32, .f64 => return error.UnsupportedOperation,
1999 }
2000 }
2001
2002 fn emitPopCount(self: *Emitter, op: *ir.Operation) EmitError!void {
2003 const result = op.getResult(0) orelse return error.InvalidArtifact;
2004 const input = try self.require(op.operands.items[0].value);
2005 switch (try scalarKind(result.type)) {
2006 .index, .u8, .u16, .u32 => {
2007 const out = self.allocU32();
2008 try self.line(" popc.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
2009 try self.bind(result, .{ .u32 = out });
2010 },
2011 .i8, .i16, .i32 => {
2012 const out = self.allocU32();
2013 try self.line(" popc.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
2014 try self.bind(result, .{ .s32 = out });
2015 },
2016 .i64, .u64 => {
2017 const out = self.allocU32();
2018 try self.line(" popc.b64 %r{d}, %rd{d};", .{ out, try self.asU64(input) });
2019 const widened = self.allocU64();
2020 try self.line(" cvt.u64.u32 %rd{d}, %r{d};", .{ widened, out });
2021 try self.bind(result, .{ .u64 = widened });
2022 },
2023 else => return error.UnsupportedOperation,
2024 }
2025 }
2026
2027 fn emitNot(self: *Emitter, op: *ir.Operation) EmitError!void {
2028 const result = op.getResult(0) orelse return error.InvalidArtifact;
2029 const input = try self.require(op.operands.items[0].value);
2030 switch (try scalarKind(result.type)) {
2031 .index, .u8, .u16, .u32 => {
2032 const out = self.allocU32();
2033 try self.line(" not.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
2034 try self.bind(result, .{ .u32 = out });
2035 },
2036 .i8, .i16, .i32 => {
2037 const out = self.allocU32();
2038 try self.line(" not.b32 %r{d}, %r{d};", .{ out, try self.asU32(input) });
2039 try self.bind(result, .{ .s32 = out });
2040 },
2041 .i64, .u64 => {
2042 const out = self.allocU64();
2043 try self.line(" not.b64 %rd{d}, %rd{d};", .{ out, try self.asU64(input) });
2044 try self.bind(result, .{ .u64 = out });
2045 },
2046 .f16, .bf16 => return error.UnsupportedOperation,
2047 .bool => {
2048 const out = self.allocPred();
2049 try self.line(" not.pred %p{d}, %p{d};", .{ out, try self.asPred(input) });
2050 try self.bind(result, .{ .pred = out });
2051 },
2052 .f32, .f64 => return error.UnsupportedOperation,
2053 }
2054 }
2055
2056 fn emitShift(self: *Emitter, op: *ir.Operation, comptime instruction: []const u8) EmitError!void {
2057 const result = op.getResult(0) orelse return error.InvalidArtifact;
2058 switch (try scalarKind(result.type)) {
2059 .index, .u8, .u16, .u32 => {
2060 const value = try self.asU32(try self.require(op.operands.items[0].value));
2061 const shift = try self.asU32(try self.require(op.operands.items[1].value));
2062 const out = self.allocU32();
2063 try self.line(" {s} %r{d}, %r{d}, %r{d};", .{ instruction, out, value, shift });
2064 try self.bind(result, .{ .u32 = out });
2065 },
2066 .i8, .i16, .i32 => {
2067 const value = try self.asU32(try self.require(op.operands.items[0].value));
2068 const shift = try self.asU32(try self.require(op.operands.items[1].value));
2069 const out = self.allocU32();
2070 try self.line(" {s} %r{d}, %r{d}, %r{d};", .{ instruction, out, value, shift });
2071 try self.bind(result, .{ .s32 = out });
2072 },
2073 .i64 => {
2074 const value = try self.asU64(try self.require(op.operands.items[0].value));
2075 const shift = try self.asU64(try self.require(op.operands.items[1].value));
2076 const amount = self.allocU32();
2077 try self.line(" cvt.u32.u64 %r{d}, %rd{d};", .{ amount, shift });
2078 const mnemonic = comptime wideShiftMnemonic(instruction, true);
2079 const out = self.allocU64();
2080 try self.line(" {s} %rd{d}, %rd{d}, %r{d};", .{ mnemonic, out, value, amount });
2081 try self.bind(result, .{ .u64 = out });
2082 },
2083 .u64 => {
2084 const value = try self.asU64(try self.require(op.operands.items[0].value));
2085 const shift = try self.asU64(try self.require(op.operands.items[1].value));
2086 const amount = self.allocU32();
2087 try self.line(" cvt.u32.u64 %r{d}, %rd{d};", .{ amount, shift });
2088 const mnemonic = comptime wideShiftMnemonic(instruction, false);
2089 const out = self.allocU64();
2090 try self.line(" {s} %rd{d}, %rd{d}, %r{d};", .{ mnemonic, out, value, amount });
2091 try self.bind(result, .{ .u64 = out });
2092 },
2093 else => return error.UnsupportedOperation,
2094 }
2095 }
2096
2097 fn wideShiftMnemonic(comptime narrow: []const u8, comptime signed: bool) []const u8 {
2098 if (std.mem.eql(u8, narrow, "shl.b32")) return "shl.b64";
2099 if (std.mem.eql(u8, narrow, "shr.s32")) return if (signed) "shr.s64" else "shr.u64";
2100 if (std.mem.eql(u8, narrow, "shr.u32")) return "shr.u64";
2101 @compileError("unsupported shift mnemonic " ++ narrow);
2102 }
2103
2104 fn emitExp(self: *Emitter, op: *ir.Operation) EmitError!void {
2105 const input = try self.asF32(try self.require(op.operands.items[0].value));
2106 const tmp = self.allocF32();
2107 const out = self.allocF32();
2108 try self.line(" mul.f32 %f{d}, %f{d}, 0f3FB8AA3B;", .{ tmp, input });
2109 try self.line(" ex2.approx.f32 %f{d}, %f{d};", .{ out, tmp });
2110 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2111 }
2112
2113 fn emitLog(self: *Emitter, op: *ir.Operation) EmitError!void {
2114 const input = try self.asF32(try self.require(op.operands.items[0].value));
2115 const tmp = self.allocF32();
2116 const out = self.allocF32();
2117 try self.line(" lg2.approx.f32 %f{d}, %f{d};", .{ tmp, input });
2118 try self.line(" mul.f32 %f{d}, %f{d}, 0f3F317218;", .{ out, tmp });
2119 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2120 }
2121
2122 fn emitVecExtract(self: *Emitter, op: *ir.Operation) EmitError!void {
2123 const extract = ArithDialect.ExtractOp{ .op = op };
2124 const vector = try self.require(extract.getVector());
2125 if (vector != .f32x4) return error.UnsupportedOperation;
2126 const lane_index = extract.getIndex() orelse return error.InvalidArtifact;
2127 if (lane_index < 0 or lane_index > 3) return error.InvalidArtifact;
2128 const lane: u32 = @intCast(lane_index);
2129 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = vector.f32x4 + lane });
2130 }
2131
2132 fn emitVecSplat(self: *Emitter, op: *ir.Operation) EmitError!void {
2133 const result = op.getResult(0) orelse return error.InvalidArtifact;
2134 if (!isVec4F32Type(result.type)) return error.UnsupportedOperation;
2135 const input = try self.asF32(try self.require(op.operands.items[0].value));
2136 const base = self.allocF32x4();
2137 var lane: u32 = 0;
2138 while (lane < 4) : (lane += 1) {
2139 try self.line(" mov.f32 %f{d}, %f{d};", .{ base + lane, input });
2140 }
2141 try self.bind(result, .{ .f32x4 = base });
2142 }
2143
2144 fn emitVecInsert(self: *Emitter, op: *ir.Operation) EmitError!void {
2145 const insert = ArithDialect.InsertOp{ .op = op };
2146 const vector = try self.require(insert.getVector());
2147 if (vector != .f32x4) return error.UnsupportedOperation;
2148 const lane_value = try self.asF32(try self.require(insert.getScalar()));
2149 const lane_index = insert.getIndex() orelse return error.InvalidArtifact;
2150 if (lane_index < 0 or lane_index > 3) return error.InvalidArtifact;
2151 const target_lane: u32 = @intCast(lane_index);
2152 const base = self.allocF32x4();
2153 var lane: u32 = 0;
2154 while (lane < 4) : (lane += 1) {
2155 const source = if (lane == target_lane) lane_value else vector.f32x4 + lane;
2156 try self.line(" mov.f32 %f{d}, %f{d};", .{ base + lane, source });
2157 }
2158 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32x4 = base });
2159 }
2160
2161 fn emitTf32Round(self: *Emitter, op: *ir.Operation) EmitError!void {
2162 const input = try self.asF32(try self.require(op.operands.items[0].value));
2163 const bits = self.allocU32();
2164 const out = self.allocF32();
2165 self.requires_sm80 = true;
2166 try self.line(" cvt.rna.tf32.f32 %r{d}, %f{d};", .{ bits, input });
2167 try self.line(" mov.b32 %f{d}, %r{d};", .{ out, bits });
2168 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2169 }
2170
2171 fn emitCpAsyncShared(self: *Emitter, op: NvptxDialect.CpAsyncSharedOp) EmitError!void {
2172 const bytes = op.getBytes() orelse return error.InvalidArtifact;
2173 if (bytes != 4 and bytes != 8 and bytes != 16) return error.UnsupportedOperation;
2174 self.requires_sm80 = true;
2175
2176 const dst = try self.require(op.getDst());
2177 const dst_info = memrefInfo(op.getDst().type) orelse return error.InvalidArtifact;
2178 if (dst_info.addr_space != .shared) return error.InvalidArtifact;
2179 const shared = switch (dst) {
2180 .shared => |shared| shared,
2181 else => return error.InvalidArtifact,
2182 };
2183 const dst_addr = try self.emitSharedAddress(shared, op.getDstIndex(), dst_info.element);
2184
2185 const src = try self.require(op.getSrc());
2186 const src_info = memrefInfo(op.getSrc().type) orelse return error.InvalidArtifact;
2187 if (!isKernelParameterAddressSpace(src_info.addr_space)) return error.InvalidArtifact;
2188 const src_ptr = switch (src) {
2189 .ptr => |ptr_reg| ptr_reg,
2190 else => return error.InvalidArtifact,
2191 };
2192 const src_addr = try self.emitPointerAddress(src_ptr, op.getSrcIndex(), src_info.element);
2193
2194 const qualifier: []const u8 = if (bytes == 16) "cg" else "ca";
2195 try self.line(" cp.async.{s}.shared.global [%r{d}+{d}], [%rd{d}+{d}], {d};", .{ qualifier, dst_addr.reg, dst_addr.imm, src_addr.reg, src_addr.imm, bytes });
2196 }
2197
2198 fn emitCpAsyncCommit(self: *Emitter, op: NvptxDialect.CpAsyncCommitOp) EmitError!void {
2199 _ = op;
2200 self.requires_sm80 = true;
2201 try self.line(" cp.async.commit_group ;", .{});
2202 }
2203
2204 fn emitCpAsyncWait(self: *Emitter, op: NvptxDialect.CpAsyncWaitOp) EmitError!void {
2205 const groups = op.getGroups() orelse return error.InvalidArtifact;
2206 self.requires_sm80 = true;
2207 try self.line(" cp.async.wait_group {d};", .{groups});
2208 }
2209
2210 fn emitMmaSync(self: *Emitter, op: NvptxDialect.MmaSyncOp) EmitError!void {
2211 const shape = op.getShape() orelse return error.InvalidArtifact;
2212 if (shape.m != 16 or shape.n != 8 or shape.k != 8) return error.UnsupportedOperation;
2213 self.requires_sm80 = true;
2214 var ab_regs: [6]u32 = undefined;
2215 for (0..ab_regs.len) |index| {
2216 const value = try self.require(op.getOperandValue(index));
2217 const bits = self.allocU32();
2218 try self.line(" mov.b32 %r{d}, %f{d};", .{ bits, try self.asF32(value) });
2219 ab_regs[index] = bits;
2220 }
2221 var c_regs: [4]u32 = undefined;
2222 for (0..c_regs.len) |index| {
2223 c_regs[index] = try self.asF32(try self.require(op.getOperandValue(6 + index)));
2224 }
2225 var d_regs: [4]u32 = undefined;
2226 for (&d_regs) |*reg| reg.* = self.allocF32();
2227 try self.line(" mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {{%f{d}, %f{d}, %f{d}, %f{d}}}, {{%r{d}, %r{d}, %r{d}, %r{d}}}, {{%r{d}, %r{d}}}, {{%f{d}, %f{d}, %f{d}, %f{d}}};", .{
2228 d_regs[0], d_regs[1], d_regs[2], d_regs[3],
2229 ab_regs[0], ab_regs[1], ab_regs[2], ab_regs[3],
2230 ab_regs[4], ab_regs[5], c_regs[0], c_regs[1],
2231 c_regs[2], c_regs[3],
2232 });
2233 for (0..d_regs.len) |index| {
2234 try self.bind(op.op.getResult(index) orelse return error.InvalidArtifact, .{ .f32 = d_regs[index] });
2235 }
2236 }
2237
2238 fn emitTanh(self: *Emitter, op: *ir.Operation) EmitError!void {
2239 const input = try self.asF32(try self.require(op.operands.items[0].value));
2240 const clamped_high = self.allocF32();
2241 const clamped = self.allocF32();
2242 const t0 = self.allocF32();
2243 const t1 = self.allocF32();
2244 const t2 = self.allocF32();
2245 const t3 = self.allocF32();
2246 const ratio = self.allocF32();
2247 const out = self.allocF32();
2248 const nan_pred = self.allocPred();
2249 try self.line(" min.f32 %f{d}, %f{d}, 0f41200000;", .{ clamped_high, input });
2250 try self.line(" max.f32 %f{d}, %f{d}, 0fC1200000;", .{ clamped, clamped_high });
2251 try self.line(" mul.f32 %f{d}, %f{d}, 0f4038AA3B;", .{ t0, clamped });
2252 try self.line(" ex2.approx.f32 %f{d}, %f{d};", .{ t1, t0 });
2253 try self.line(" add.f32 %f{d}, %f{d}, 0f3F800000;", .{ t2, t1 });
2254 try self.line(" sub.f32 %f{d}, %f{d}, 0f3F800000;", .{ t3, t1 });
2255 try self.line(" div.approx.f32 %f{d}, %f{d}, %f{d};", .{ ratio, t3, t2 });
2256 try self.line(" testp.notanumber.f32 %p{d}, %f{d};", .{ nan_pred, input });
2257 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ out, input, ratio, nan_pred });
2258 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2259 }
2260
2261 fn emitTan(self: *Emitter, op: *ir.Operation) EmitError!void {
2262 const input = try self.asF32(try self.require(op.operands.items[0].value));
2263 const s = self.allocF32();
2264 const c = self.allocF32();
2265 const out = self.allocF32();
2266 try self.line(" sin.approx.f32 %f{d}, %f{d};", .{ s, input });
2267 try self.line(" cos.approx.f32 %f{d}, %f{d};", .{ c, input });
2268 try self.line(" div.approx.f32 %f{d}, %f{d}, %f{d};", .{ out, s, c });
2269 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2270 }
2271
2272 fn emitPow(self: *Emitter, op: *ir.Operation) EmitError!void {
2273 const base = try self.asF32(try self.require(op.operands.items[0].value));
2274 const exponent = try self.asF32(try self.require(op.operands.items[1].value));
2275 const log_base = self.allocF32();
2276 const scaled = self.allocF32();
2277 const out = self.allocF32();
2278 try self.line(" lg2.approx.f32 %f{d}, %f{d};", .{ log_base, base });
2279 try self.line(" mul.f32 %f{d}, %f{d}, %f{d};", .{ scaled, log_base, exponent });
2280 try self.line(" ex2.approx.f32 %f{d}, %f{d};", .{ out, scaled });
2281 try self.bind(op.getResult(0) orelse return error.InvalidArtifact, .{ .f32 = out });
2282 }
2283
2284 fn emitAtan2(self: *Emitter, op: *ir.Operation) EmitError!void {
2285 const result = op.getResult(0) orelse return error.InvalidArtifact;
2286 if ((try computeKind(result.type)) != .f32) return error.UnsupportedOperation;
2287 const y = try self.asF32(try self.require(op.operands.items[0].value));
2288 const x = try self.asF32(try self.require(op.operands.items[1].value));
2289 const abs_y = self.allocF32();
2290 const abs_x = self.allocF32();
2291 const y_over_x = self.allocF32();
2292 const x_over_y = self.allocF32();
2293 const ratio = self.allocF32();
2294 const complement = self.allocF32();
2295 const base = self.allocF32();
2296 const zero = self.allocF32();
2297 const half_pi = self.allocF32();
2298 const pi = self.allocF32();
2299 const x_axis = self.allocF32();
2300 const axis = self.allocF32();
2301 const pi_minus = self.allocF32();
2302 const quadrant = self.allocF32();
2303 const negated = self.allocF32();
2304 const out = self.allocF32();
2305 const use_y_over_x = self.allocPred();
2306 const x_zero = self.allocPred();
2307 const y_zero = self.allocPred();
2308 const x_negative = self.allocPred();
2309 const y_negative = self.allocPred();
2310 try self.line(" abs.f32 %f{d}, %f{d};", .{ abs_y, y });
2311 try self.line(" abs.f32 %f{d}, %f{d};", .{ abs_x, x });
2312 try self.line(" div.approx.f32 %f{d}, %f{d}, %f{d};", .{ y_over_x, abs_y, abs_x });
2313 try self.line(" div.approx.f32 %f{d}, %f{d}, %f{d};", .{ x_over_y, abs_x, abs_y });
2314 try self.line(" setp.gt.f32 %p{d}, %f{d}, %f{d};", .{ use_y_over_x, abs_x, abs_y });
2315 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ ratio, y_over_x, x_over_y, use_y_over_x });
2316 const atan = try self.emitAtanApprox(ratio);
2317 try self.line(" mov.f32 %f{d}, 0f00000000;", .{zero});
2318 try self.line(" mov.f32 %f{d}, 0f3FC90FDB;", .{half_pi});
2319 try self.line(" mov.f32 %f{d}, 0f40490FDB;", .{pi});
2320 try self.line(" sub.f32 %f{d}, %f{d}, %f{d};", .{ complement, half_pi, atan });
2321 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ base, atan, complement, use_y_over_x });
2322 try self.line(" setp.eq.f32 %p{d}, %f{d}, %f{d};", .{ x_zero, abs_x, zero });
2323 try self.line(" setp.eq.f32 %p{d}, %f{d}, %f{d};", .{ y_zero, abs_y, zero });
2324 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ x_axis, half_pi, base, x_zero });
2325 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ axis, zero, x_axis, y_zero });
2326 try self.line(" setp.lt.f32 %p{d}, %f{d}, %f{d};", .{ x_negative, x, zero });
2327 try self.line(" sub.f32 %f{d}, %f{d}, %f{d};", .{ pi_minus, pi, axis });
2328 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ quadrant, pi_minus, axis, x_negative });
2329 try self.line(" setp.lt.f32 %p{d}, %f{d}, %f{d};", .{ y_negative, y, zero });
2330 try self.line(" neg.f32 %f{d}, %f{d};", .{ negated, quadrant });
2331 try self.line(" selp.f32 %f{d}, %f{d}, %f{d}, %p{d};", .{ out, negated, quadrant, y_negative });
2332 try self.bind(result, .{ .f32 = out });
2333 }
2334
2335 fn emitAtanApprox(self: *Emitter, input: u32) EmitError!u32 {
2336 const abs_value = self.allocF32();
2337 const centered = self.allocF32();
2338 const b_term = self.allocF32();
2339 const coeff = self.allocF32();
2340 const term0 = self.allocF32();
2341 const term1 = self.allocF32();
2342 const linear = self.allocF32();
2343 const out = self.allocF32();
2344 try self.line(" abs.f32 %f{d}, %f{d};", .{ abs_value, input });
2345 try self.line(" sub.f32 %f{d}, %f{d}, 0f3F800000;", .{ centered, abs_value });
2346 try self.line(" mul.f32 %f{d}, %f{d}, 0f3D87C84B;", .{ b_term, abs_value });
2347 try self.line(" add.f32 %f{d}, %f{d}, 0f3E7A92A3;", .{ coeff, b_term });
2348 try self.line(" mul.f32 %f{d}, %f{d}, %f{d};", .{ term0, input, centered });
2349 try self.line(" mul.f32 %f{d}, %f{d}, %f{d};", .{ term1, term0, coeff });
2350 try self.line(" mul.f32 %f{d}, %f{d}, 0f3F490FDB;", .{ linear, input });
2351 try self.line(" sub.f32 %f{d}, %f{d}, %f{d};", .{ out, linear, term1 });
2352 return out;
2353 }
2354
2355 fn emitFma(self: *Emitter, op: *ir.Operation) EmitError!void {
2356 const result = op.getResult(0) orelse return error.InvalidArtifact;
2357 if ((try computeKind(result.type)) != .f32) return error.UnsupportedOperation;
2358 const a = try self.asF32(try self.require(op.operands.items[0].value));
2359 const b = try self.asF32(try self.require(op.operands.items[1].value));
2360 const c = try self.asF32(try self.require(op.operands.items[2].value));
2361 const out = self.allocF32();
2362 try self.line(" fma.rn.f32 %f{d}, %f{d}, %f{d}, %f{d};", .{ out, a, b, c });
2363 try self.bind(result, .{ .f32 = out });
2364 }
2365
2366 fn emitIf(self: *Emitter, op: ScfDialect.IfOp) EmitError!void {
2367 const result_count = op.getNumResults();
2368 if (result_count == 0) {
2369 const cond = try self.asPred(try self.require(op.getCondition()));
2370 const id = self.freshLabel();
2371 try self.line(" @!%p{d} bra LIF_ELSE_{d};", .{ cond, id });
2372 try self.emitBlock(op.getThenBlock());
2373 try self.line(" bra LIF_DONE_{d};", .{id});
2374 try self.line("LIF_ELSE_{d}:", .{id});
2375 if (op.getElseBlock()) |else_block| {
2376 try self.emitBlock(else_block);
2377 }
2378 try self.line("LIF_DONE_{d}:", .{id});
2379 return;
2380 }
2381
2382 const else_block = op.getElseBlock() orelse return error.InvalidArtifact;
2383 const cond = try self.asPred(try self.require(op.getCondition()));
2384 const id = self.freshLabel();
2385
2386 var result_registers: std.ArrayListUnmanaged(Value) = .empty;
2387 defer result_registers.deinit(self.allocator);
2388
2389 try self.line(" @!%p{d} bra LIF_ELSE_{d};", .{ cond, id });
2390 try self.emitYieldingBlock(op.getThenBlock(), result_count, &result_registers);
2391 try self.line(" bra LIF_DONE_{d};", .{id});
2392 try self.line("LIF_ELSE_{d}:", .{id});
2393 try self.emitYieldingBlock(else_block, result_count, &result_registers);
2394 try self.line("LIF_DONE_{d}:", .{id});
2395 self.resetAddressBases();
2396
2397 for (0..result_count) |index| {
2398 const result = op.op.getResult(index) orelse return error.InvalidArtifact;
2399 try self.bind(result, result_registers.items[index]);
2400 }
2401 }
2402
2403 fn emitYieldingBlock(
2404 self: *Emitter,
2405 block: *ir.Block,
2406 result_count: usize,
2407 result_registers: *std.ArrayListUnmanaged(Value),
2408 ) EmitError!void {
2409 self.resetAddressBases();
2410 var ops = block.getOperations();
2411 while (ops.next()) |block_op| {
2412 if (std.mem.eql(u8, block_op.name.name, ScfDialect.YieldOp.operation_name)) {
2413 const yield = ScfDialect.YieldOp{ .op = block_op };
2414 const yielded = yield.getOperands();
2415 if (yielded.len != result_count) return error.InvalidArtifact;
2416 for (yielded, 0..) |yield_value, index| {
2417 const source = try self.require(yield_value);
2418 if (result_registers.items.len <= index) {
2419 const register = try self.allocLike(source);
2420 result_registers.append(self.allocator, register) catch return error.OutOfMemory;
2421 }
2422 try self.emitMove(result_registers.items[index], source);
2423 }
2424 return;
2425 }
2426 try self.emitOperation(block_op);
2427 }
2428 return error.InvalidArtifact;
2429 }
2430
2431 fn emitWhile(self: *Emitter, op: ScfDialect.WhileOp) EmitError!void {
2432 const before = op.getBeforeBlock();
2433 const after = op.getAfterBlock();
2434 const carry_count = op.op.operands.items.len;
2435 if (op.op.results.items.len != carry_count) return error.InvalidArtifact;
2436 if (before.arguments.items.len != carry_count) return error.InvalidArtifact;
2437 if (after.arguments.items.len != carry_count) return error.InvalidArtifact;
2438
2439 var carries: std.ArrayListUnmanaged(Value) = .empty;
2440 defer carries.deinit(self.allocator);
2441 for (op.op.operands.items, before.arguments.items) |operand, before_arg| {
2442 const initial = try self.require(operand.value);
2443 const carry = try self.allocLike(initial);
2444 try self.emitMove(carry, initial);
2445 carries.append(self.allocator, carry) catch return error.OutOfMemory;
2446 try self.bind(before_arg, carry);
2447 }
2448
2449 var exits: std.ArrayListUnmanaged(Value) = .empty;
2450 defer exits.deinit(self.allocator);
2451
2452 const id = self.freshLabel();
2453 self.resetAddressBases();
2454 try self.line("LWHILE_HEAD_{d}:", .{id});
2455
2456 var before_ops = before.getOperations();
2457 var saw_condition = false;
2458 while (before_ops.next()) |before_op| {
2459 if (std.mem.eql(u8, before_op.name.name, ScfDialect.ConditionOp.operation_name)) {
2460 const condition = ScfDialect.ConditionOp{ .op = before_op };
2461 const args = condition.getArgs();
2462 if (args.len != carry_count) return error.InvalidArtifact;
2463 for (args, 0..) |arg, index| {
2464 const source = try self.require(arg);
2465 if (exits.items.len <= index) {
2466 const register = try self.allocLike(source);
2467 exits.append(self.allocator, register) catch return error.OutOfMemory;
2468 }
2469 try self.emitMove(exits.items[index], source);
2470 }
2471 const cond = try self.asPred(try self.require(condition.getCondition()));
2472 try self.line(" @!%p{d} bra LWHILE_DONE_{d};", .{ cond, id });
2473 saw_condition = true;
2474 break;
2475 }
2476 try self.emitOperation(before_op);
2477 }
2478 if (!saw_condition) return error.InvalidArtifact;
2479
2480 for (after.arguments.items, 0..) |after_arg, index| {
2481 try self.bind(after_arg, exits.items[index]);
2482 }
2483
2484 var after_ops = after.getOperations();
2485 var saw_yield = false;
2486 while (after_ops.next()) |after_op| {
2487 if (std.mem.eql(u8, after_op.name.name, ScfDialect.YieldOp.operation_name)) {
2488 const yield = ScfDialect.YieldOp{ .op = after_op };
2489 const yielded = yield.getOperands();
2490 if (yielded.len != carry_count) return error.InvalidArtifact;
2491 for (yielded, carries.items) |yield_value, carry| {
2492 try self.emitMove(carry, try self.require(yield_value));
2493 }
2494 saw_yield = true;
2495 break;
2496 }
2497 try self.emitOperation(after_op);
2498 }
2499 if (!saw_yield) return error.InvalidArtifact;
2500
2501 try self.line(" bra LWHILE_HEAD_{d};", .{id});
2502 try self.line("LWHILE_DONE_{d}:", .{id});
2503 self.resetAddressBases();
2504
2505 for (op.op.results.items, exits.items) |*result, exit| {
2506 try self.bind(result, exit);
2507 }
2508 }
2509
2510 fn emitFor(self: *Emitter, op: ScfDialect.ForOp) EmitError!void {
2511 const integer = try LoopInteger.fromType(op.getLowerBound().type);
2512 const lower = try self.require(op.getLowerBound());
2513 const upper = try integer.register(self, try self.require(op.getUpperBound()));
2514 const step = try integer.register(self, try self.require(op.getStep()));
2515 const body_block = op.getBodyBlock();
2516 const iter_args = op.getInitArgs();
2517 if (iter_args.len != op.op.results.items.len) return error.InvalidArtifact;
2518 if (body_block.arguments.items.len != iter_args.len + 1) return error.InvalidArtifact;
2519
2520 const induction = try self.allocLike(lower);
2521 const iv = try integer.register(self, induction);
2522 try self.emitMove(induction, lower);
2523 try self.bind(op.getInductionVar(), induction);
2524
2525 const accumulators = try self.allocator.alloc(Value, iter_args.len);
2526 defer self.allocator.free(accumulators);
2527 for (iter_args, body_block.arguments.items[1..], accumulators) |init_arg, block_arg, *acc| {
2528 const initial = try self.require(init_arg);
2529 acc.* = try self.allocLike(initial);
2530 try self.emitMove(acc.*, initial);
2531 try self.bind(block_arg, acc.*);
2532 }
2533
2534 const pred = self.allocPred();
2535 const id = self.freshLabel();
2536 const registers = integer.registerFile();
2537 self.resetAddressBases();
2538 try self.line("LFOR_HEAD_{d}:", .{id});
2539 try self.line(" setp.ge.{s} %p{d}, %{s}{d}, %{s}{d};", .{
2540 @tagName(integer), pred, registers, iv, registers, upper,
2541 });
2542 try self.line(" @%p{d} bra LFOR_DONE_{d};", .{ pred, id });
2543
2544 try self.emitLoopBody(body_block, accumulators);
2545 self.resetAddressBases();
2546 try self.line(" add.{s} %{s}{d}, %{s}{d}, %{s}{d};", .{
2547 @tagName(integer), registers, iv, registers, iv, registers, step,
2548 });
2549 try self.line(" bra LFOR_HEAD_{d};", .{id});
2550 try self.line("LFOR_DONE_{d}:", .{id});
2551
2552 for (op.op.results.items, accumulators) |*result, acc| {
2553 try self.bind(result, acc);
2554 }
2555 }
2556
2557 fn emitLoopBody(self: *Emitter, body_block: *ir.Block, accumulators: []const Value) EmitError!void {
2558 var ops = body_block.getOperations();
2559 while (ops.next()) |body_op| {
2560 if (std.mem.eql(u8, body_op.name.name, ScfDialect.YieldOp.operation_name)) {
2561 const yield = ScfDialect.YieldOp{ .op = body_op };
2562 try self.emitParallelMoves(accumulators, yield.getOperands());
2563 return;
2564 }
2565 try self.emitOperation(body_op);
2566 }
2567 return error.InvalidArtifact;
2568 }
2569
2570 fn emitParallelMoves(
2571 self: *Emitter,
2572 destinations: []const Value,
2573 sources: []const *ir.Value,
2574 ) EmitError!void {
2575 if (destinations.len != sources.len) return error.InvalidArtifact;
2576 const snapshots = try self.allocator.alloc(Value, sources.len);
2577 defer self.allocator.free(snapshots);
2578 for (sources, snapshots) |source, *snapshot| {
2579 const value = try self.require(source);
2580 snapshot.* = try self.allocLike(value);
2581 try self.emitMove(snapshot.*, value);
2582 }
2583 for (destinations, snapshots) |destination, snapshot| {
2584 try self.emitMove(destination, snapshot);
2585 }
2586 }
2587
2588 fn allocLike(self: *Emitter, value: Value) abi.Error!Value {
2589 return switch (value) {
2590 .pred => .{ .pred = self.allocPred() },
2591 .u32 => .{ .u32 = self.allocU32() },
2592 .s32 => .{ .s32 = self.allocU32() },
2593 .u64 => .{ .u64 = self.allocU64() },
2594 .f32 => .{ .f32 = self.allocF32() },
2595 .f32x4 => .{ .f32x4 = self.allocF32x4() },
2596 .f64 => .{ .f64 = self.allocF64() },
2597 .ptr, .shared => error.UnsupportedOperation,
2598 };
2599 }
2600
2601 fn emitMove(self: *Emitter, dst: Value, src: Value) EmitError!void {
2602 switch (dst) {
2603 .pred => |reg| try self.line(" mov.pred %p{d}, %p{d};", .{ reg, try self.asPred(src) }),
2604 .u32, .s32 => |reg| try self.line(" mov.u32 %r{d}, %r{d};", .{ reg, try self.asU32(src) }),
2605 .u64 => |reg| try self.line(" mov.u64 %rd{d}, %rd{d};", .{ reg, try self.asU64(src) }),
2606 .f32 => |reg| try self.line(" mov.f32 %f{d}, %f{d};", .{ reg, try self.asF32(src) }),
2607 .f32x4 => |reg| {
2608 if (src != .f32x4) return error.UnsupportedOperation;
2609 var lane: u32 = 0;
2610 while (lane < 4) : (lane += 1) {
2611 try self.line(" mov.f32 %f{d}, %f{d};", .{ reg + lane, src.f32x4 + lane });
2612 }
2613 },
2614 .f64 => |reg| try self.line(" mov.f64 %fd{d}, %fd{d};", .{ reg, try self.asF64(src) }),
2615 .ptr, .shared => return error.UnsupportedOperation,
2616 }
2617 }
2618
2619 fn bind(self: *Emitter, value: *ir.Value, ptx_value: Value) abi.Error!void {
2620 self.values.put(self.allocator, value, ptx_value) catch return error.OutOfMemory;
2621 }
2622
2623 fn require(self: *Emitter, value: *ir.Value) abi.Error!Value {
2624 return self.values.get(value) orelse error.InvalidArtifact;
2625 }
2626
2627 fn asU32(_: *Emitter, value: Value) abi.Error!u32 {
2628 return switch (value) {
2629 .u32, .s32 => |reg| reg,
2630 else => error.InvalidArtifact,
2631 };
2632 }
2633
2634 fn asF32(_: *Emitter, value: Value) abi.Error!u32 {
2635 return switch (value) {
2636 .f32 => |reg| reg,
2637 else => error.InvalidArtifact,
2638 };
2639 }
2640
2641 fn asF64(_: *Emitter, value: Value) abi.Error!u32 {
2642 return switch (value) {
2643 .f64 => |reg| reg,
2644 else => error.InvalidArtifact,
2645 };
2646 }
2647
2648 fn asU64(_: *Emitter, value: Value) abi.Error!u32 {
2649 return switch (value) {
2650 .u64 => |reg| reg,
2651 else => error.InvalidArtifact,
2652 };
2653 }
2654
2655 fn asPred(_: *Emitter, value: Value) abi.Error!u32 {
2656 return switch (value) {
2657 .pred => |reg| reg,
2658 else => error.InvalidArtifact,
2659 };
2660 }
2661
2662 fn allocU32(self: *Emitter) u32 {
2663 const reg = self.next_r;
2664 self.next_r += 1;
2665 return reg;
2666 }
2667
2668 fn allocB16(self: *Emitter) u32 {
2669 const reg = self.next_h;
2670 self.next_h += 1;
2671 return reg;
2672 }
2673
2674 fn allocF32(self: *Emitter) u32 {
2675 const reg = self.next_f;
2676 self.next_f += 1;
2677 return reg;
2678 }
2679
2680 fn allocF32x4(self: *Emitter) u32 {
2681 const reg = self.next_f;
2682 self.next_f += 4;
2683 return reg;
2684 }
2685
2686 fn allocF64(self: *Emitter) u32 {
2687 const reg = self.next_fd;
2688 self.next_fd += 1;
2689 return reg;
2690 }
2691
2692 fn allocPtr(self: *Emitter) u32 {
2693 const reg = self.next_rd;
2694 self.next_rd += 1;
2695 return reg;
2696 }
2697
2698 fn allocU64(self: *Emitter) u32 {
2699 const reg = self.next_rd;
2700 self.next_rd += 1;
2701 return reg;
2702 }
2703
2704 fn allocPred(self: *Emitter) u32 {
2705 const reg = self.next_p;
2706 self.next_p += 1;
2707 return reg;
2708 }
2709
2710 fn allocShared(self: *Emitter) u32 {
2711 const id = self.next_shared;
2712 self.next_shared += 1;
2713 return id;
2714 }
2715
2716 fn emitBf16BitsToF32(self: *Emitter, bits: u32) EmitError!u32 {
2717 const widened = self.allocU32();
2718 const out = self.allocF32();
2719 try self.line(" shl.b32 %r{d}, %r{d}, 16;", .{ widened, bits });
2720 try self.line(" mov.b32 %f{d}, %r{d};", .{ out, widened });
2721 return out;
2722 }
2723
2724 fn emitF32ToBf16Bits(self: *Emitter, value: Value) EmitError!u32 {
2725 const bits = self.allocU32();
2726 const out = self.allocU32();
2727 try self.line(" mov.b32 %r{d}, %f{d};", .{ bits, try self.asF32(value) });
2728 try self.line(" shr.u32 %r{d}, %r{d}, 16;", .{ out, bits });
2729 return out;
2730 }
2731
2732 fn freshLabel(self: *Emitter) u32 {
2733 const label = self.next_label;
2734 self.next_label += 1;
2735 return label;
2736 }
2737
2738 fn line(self: *Emitter, comptime fmt: []const u8, args: anytype) std.Io.Writer.Error!void {
2739 try self.body.writer.print(fmt, args);
2740 try self.body.writer.writeByte('\n');
2741 }
2742 };
2743
2744 const MemrefInfo = struct {
2745 size: ?u64,
2746 element: ScalarKind,
2747 addr_space: dialects.AddressSpace,
2748 alignment: ?u64,
2749 };
2750
2751 fn writeHeader(writer: *std.Io.Writer, requires_sm80: bool) std.Io.Writer.Error!void {
2752 if (requires_sm80) {
2753 try writer.writeAll(".version 7.0\n.target sm_80\n.address_size 64\n\n");
2754 return;
2755 }
2756 try writer.writeAll(".version 6.2\n.target sm_52\n.address_size 64\n\n");
2757 }
2758
2759 fn scalarKind(typ: ir.Type) abi.Error!ScalarKind {
2760 return scalar_kinds.kindFromType(typ) orelse error.UnsupportedOperation;
2761 }
2762
2763 fn isVec4F32Type(typ: ir.Type) bool {
2764 const name = typ.getDialectTypeName() orelse return false;
2765 return std.mem.eql(u8, name, arith_names.vec4xf32);
2766 }
2767
2768 fn computeKind(typ: ir.Type) abi.Error!ScalarKind {
2769 const kind = try scalarKind(typ);
2770 return switch (kind) {
2771 .i8, .i16 => .i32,
2772 .u8, .u16 => .u32,
2773 .f16, .bf16 => .f32,
2774 else => kind,
2775 };
2776 }
2777
2778 fn memrefInfo(typ: ir.Type) ?MemrefInfo {
2779 const name = typ.getDialectTypeName() orelse return null;
2780 if (!std.mem.eql(u8, name, MemrefDialect.name)) return null;
2781 const params = MemrefDialect.parseMemrefParams(typ.getDialectParamKey() orelse return null) orelse return null;
2782 return .{
2783 .size = params.size,
2784 .element = scalar_kinds.kindFromTypeName(params.element_type_name) orelse return null,
2785 .addr_space = params.addr_space,
2786 .alignment = params.alignment,
2787 };
2788 }
2789
2790 fn dynamicSharedByteOffset(op: *ir.Operation) abi.Error!u32 {
2791 const attr = op.getAttrAs(ir.Attribute.IntegerAttr, gpu.attr_names.dynamic_shared_byte_offset) orelse return 0;
2792 const value = attr.getValue();
2793 if (value < 0) return error.InvalidArtifact;
2794 return std.math.cast(u32, value) orelse return error.InvalidArtifact;
2795 }
2796
2797 fn isKernelParameterAddressSpace(addr_space: dialects.AddressSpace) bool {
2798 return switch (addr_space) {
2799 .host, .device, .constant, .unified => true,
2800 .shared, .local => false,
2801 };
2802 }
2803
2804 fn elementByteSize(kind: ScalarKind) u32 {
2805 return switch (kind) {
2806 .bool => 1,
2807 .i8, .u8 => 1,
2808 .i16, .u16 => 2,
2809 .i64, .u64 => 8,
2810 .f64 => 8,
2811 .f16 => 2,
2812 .bf16 => 2,
2813 .index, .i32, .u32, .f32 => 4,
2814 };
2815 }
2816
2817 fn dimName(dim: gpu.Dimension) []const u8 {
2818 return switch (dim) {
2819 .x => "x",
2820 .y => "y",
2821 .z => "z",
2822 };
2823 }
2824
2825 fn ptxPredicate(pred: CmpPredicate) []const u8 {
2826 return switch (pred) {
2827 .eq => "eq",
2828 .ne => "ne",
2829 .lt, .slt, .ult => "lt",
2830 .le, .sle, .ule => "le",
2831 .gt, .sgt, .ugt => "gt",
2832 .ge, .sge, .uge => "ge",
2833 };
2834 }
2835
2836 fn ptxShuffleMode(mode: gpu.ShuffleMode) []const u8 {
2837 return switch (mode) {
2838 .sync => "idx",
2839 .down => "down",
2840 .up => "up",
2841 .xor => "bfly",
2842 };
2843 }
2844
2845 const warp_reduce_deltas = [_]u32{ 16, 8, 4, 2, 1 };
2846 const warp_scan_offsets = [_]u32{ 1, 2, 4, 8, 16 };
2847
2848 fn f32Bits(value: f64) u32 {
2849 const narrowed: f32 = @floatCast(value);
2850 return @bitCast(narrowed);
2851 }
2852
2853 fn f64Bits(value: f64) u64 {
2854 return @bitCast(value);
2855 }
2856
2857 test "cuda scalar support mask follows Choir scalar spellings" {
2858 inline for (std.meta.tags(ScalarKind)) |kind| {
2859 try std.testing.expectEqual(
2860 @as(?ScalarKind, kind),
2861 scalar_kinds.kindFromTypeName(dialects.arith.scalarTypeName(kind)),
2862 );
2863 }
2864 try std.testing.expectEqual(@as(?ScalarKind, null), scalar_kinds.kindFromTypeName("arith.unknown"));
2865 }