aot.zig: a launch with no exact captured variant takes one that assumes less of its arguments

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,23 +195,90 @@ 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;
                 if (!std.mem.eql(u8, p.type, "fp32")) return false;
             },
             .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 = &params,
+        .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);
+}