diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index b8e58d2e8e..9f1dcf46b8 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -122,7 +122,12 @@ pub const Set = struct { pub fn find(self: *const Set, function: []const u8, args: []const Arg, consts: []const Const) !*const Variant { for (self.variants) |*v| { if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; - if (matches(v.spec, args, consts)) return v; + if (matches(v.spec, args, consts, .exact)) return v; + } + // a variant that assumes less of an argument than the launch gives computes the same for it, so it may serve + for (self.variants) |*v| { + if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; + if (matches(v.spec, args, consts, .general)) return v; } std.log.err("no captured Triton variant of {s} for this launch:", .{function}); for (args) |a| switch (a.value) { @@ -167,7 +172,9 @@ fn lookup(args: []const Arg, name: []const u8) ?Arg { return null; } -fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { +const Fit = enum { exact, general }; + +fn matches(k: KernelJson, args: []const Arg, consts: []const Const, fit: Fit) bool { for (consts) |c| { const got = k.consts.map.get(c.name) orelse return false; if (c.int) |x| if (got.int == null or got.int.? != x) return false; @@ -188,12 +195,20 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { } const p = param.?; if (!std.mem.eql(u8, p.type, "i32")) return false; - if (!p.nospec and x == 1) return false; - if (p.div16 != (!p.nospec and @mod(x, 16) == 0)) return false; + const div = !p.nospec and @mod(x, 16) == 0; + switch (fit) { + .exact => if ((!p.nospec and x == 1) or p.div16 != div) return false, + .general => if (p.div16 and !div) return false, + } }, .ptr => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, x.ty) or p.div16 != (x.addr % 16 == 0)) return false; + if (!std.mem.eql(u8, p.type, x.ty)) return false; + const div = x.addr % 16 == 0; + switch (fit) { + .exact => if (p.div16 != div) return false, + .general => if (p.div16 and !div) return false, + } }, .f32 => { const p = param orelse return false; @@ -201,10 +216,69 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { }, .u64 => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, "u64") or p.div16 != (x % 16 == 0)) return false; + if (!std.mem.eql(u8, p.type, "u64")) return false; + const div = x % 16 == 0; + switch (fit) { + .exact => if (p.div16 != div) return false, + .general => if (p.div16 and !div) return false, + } }, } runtime += 1; } return runtime == k.params.len; } + +/// A set parsed from aot.json text whose variants carry no loaded cubin, so find can be tested without a GPU. +fn testSet(gpa: std.mem.Allocator, text: []const u8) !Set { + const parsed = try std.json.parseFromSlice(SetJson, gpa, text, .{ .ignore_unknown_fields = true, .allocate = .alloc_always }); + errdefer parsed.deinit(); + const variants = try gpa.alloc(Variant, parsed.value.kernels.len); + for (parsed.value.kernels, variants) |k, *v| v.* = .{ .spec = k, .kernel = undefined }; + return .{ .parsed = parsed, .variants = variants, .gpa = gpa }; +} + +/// Frees a set from testSet: its variants never loaded a cubin, so no module is unloaded. +fn freeTestSet(gpa: std.mem.Allocator, set: *Set) void { + gpa.free(set.variants); + set.parsed.deinit(); +} + +test "a variant captured without div16 serves a 16-aligned pointer" { + const gpa = std.testing.allocator; + var set = try testSet(gpa, + \\{"kernels":[{"fn":"k","hash":"a","name":"k","num_warps":4, + \\ "params":[{"name":"P","type":"*i8","div16":false,"nospec":false}],"consts":{}}]} + ); + defer freeTestSet(gpa, &set); + const v = try set.find("k", &.{ptr("P", "*i8", 0x1000)}, &.{}); + try std.testing.expectEqualStrings("a", v.spec.hash); +} + +test "a variant captured with div16 refuses an unaligned pointer" { + var params = [_]ParamJson{.{ .name = "P", .type = "*i8", .div16 = true, .nospec = false }}; + const k: KernelJson = .{ + .@"fn" = "k", + .hash = "a", + .name = "k", + .num_warps = 4, + .params = ¶ms, + .consts = .{}, + }; + const unaligned = &[_]Arg{ptr("P", "*i8", 0x1008)}; + try std.testing.expect(!matches(k, unaligned, &.{}, .exact)); + try std.testing.expect(!matches(k, unaligned, &.{}, .general)); +} + +test "an aligned pointer takes the div16 variant when the set holds both" { + const gpa = std.testing.allocator; + var set = try testSet(gpa, + \\{"kernels":[{"fn":"k","hash":"a","name":"k","num_warps":4, + \\ "params":[{"name":"P","type":"*i8","div16":false,"nospec":false}],"consts":{}}, + \\ {"fn":"k","hash":"b","name":"k","num_warps":4, + \\ "params":[{"name":"P","type":"*i8","div16":true,"nospec":false}],"consts":{}}]} + ); + defer freeTestSet(gpa, &set); + const v = try set.find("k", &.{ptr("P", "*i8", 0x1000)}, &.{}); + try std.testing.expectEqualStrings("b", v.spec.hash); +}