lib/accy/src/preparation/numerics.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const gpu = @import("gpu");
   3 
   4 const accy_choir = @import("../choir/root.zig");
   5 const executable = @import("../executable/root.zig");
   6 const kernel = @import("../kernel/root.zig");
   7 const kernelization = @import("kernelization/root.zig");
   8 const pipeline = @import("pipeline.zig");
   9 
  10 const tensor = @import("../tensor/root.zig");
  11 const SemanticBuilder = accy_choir.SemanticBuilder;
  12 const SemanticModule = accy_choir.SemanticModule;
  13 
  14 const element_count = 8;
  15 
  16 fn runGeneratedProgram(
  17     allocator: std.mem.Allocator,
  18     program: *kernelization.GeneratedKernelProgram,
  19     args: []const kernel.Argument,
  20 ) !void {
  21     var executor = kernel.program.execution.Executor.init(program);
  22     defer executor.deinit();
  23     try executor.runCpu(allocator, args);
  24 }
  25 
  26 fn compileAndLoadSemanticModule(
  27     allocator: std.mem.Allocator,
  28     handle: gpu.BackendHandle,
  29     module: *SemanticModule,
  30     options: executable.FragmentCompilerOptions,
  31 ) !*executable.LoadedFragment {
  32     const compiled = try executable.compileFragmentFromSemanticModule(allocator, handle, module, options);
  33     return try executable.loadFragment(allocator, handle, compiled, options);
  34 }
  35 
  36 fn runSingleKernel(
  37     allocator: std.mem.Allocator,
  38     module: *SemanticModule,
  39     args: []const kernel.Argument,
  40 ) !void {
  41     const contract_module = try pipeline.prepareContractJobFromSemanticModule(allocator, module, .{});
  42     const tensor_module = try pipeline.prepareTensorJobFromContractJob(allocator, contract_module, .{});
  43     const dispatch_module = try pipeline.prepareDispatchJobFromTensorJob(allocator, tensor_module, .{});
  44     const memory_module = try pipeline.prepareMemoryJobFromDispatchJob(allocator, dispatch_module, .{});
  45     defer memory_module.deinit();
  46 
  47     var pass_ctx = memory_module.passContext();
  48     defer pass_ctx.deinit();
  49     const analysis = try kernelization.getKernelizationAnalysis(&pass_ctx, memory_module.choir_module);
  50     try std.testing.expectEqual(@as(usize, 1), analysis.kernelCount());
  51     try runGeneratedProgram(allocator, &analysis.kernels.items[0].program, args);
  52 }
  53 
  54 fn binaryF32Module(allocator: std.mem.Allocator, comptime op: enum { add, mul }) !*SemanticModule {
  55     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
  56     errdefer builder.deinit();
  57     const vec_ty = try builder.tensor(.f32, &.{element_count});
  58     var fb = try builder.beginFunction("numerics_" ++ @tagName(op), &.{ vec_ty, vec_ty }, &.{vec_ty});
  59     const out = switch (op) {
  60         .add => try fb.add(fb.parameter(0), fb.parameter(1)),
  61         .mul => try fb.mul(fb.parameter(0), fb.parameter(1)),
  62     };
  63     try fb.return_(&.{out});
  64     try fb.finish();
  65     return try builder.finish();
  66 }
  67 
  68 fn cumsumModule(allocator: std.mem.Allocator, comptime total: i64) !*SemanticModule {
  69     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
  70     errdefer builder.deinit();
  71     const vec_ty = try builder.tensor(.f32, &.{total});
  72     var fb = try builder.beginFunction("numerics_cumsum_schedule", &.{vec_ty}, &.{vec_ty});
  73     const out_value = try fb.cumsum(fb.parameter(0), vec_ty, 0);
  74     try fb.return_(&.{out_value});
  75     try fb.finish();
  76     return try builder.finish();
  77 }
  78 
  79 fn expectCumsumGeneratedSchedule(comptime total: i64, expected_threads: u32, expected_blocks: u32) !void {
  80     const allocator = std.testing.allocator;
  81     var state = gpu.recording.BackendState{
  82         .allocator = allocator,
  83         .kind = .cuda,
  84         .format = .cuda_ptx,
  85     };
  86 
  87     const module = try cumsumModule(allocator, total);
  88     var prepared = try executable.prepareFragmentFromSemanticModule(allocator, state.handle(), module, .{});
  89     defer prepared.deinit();
  90 
  91     try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
  92     const summary = try prepared.generatedKernelSummary(0);
  93     try std.testing.expectEqual(kernelization.GeneratedScheduleKind.flat, summary.schedule.kind);
  94     try std.testing.expectEqual(expected_threads, summary.schedule.threads.x);
  95     var entry_buffer: [128]u8 = undefined;
  96     const expected_entry_part = try std.fmt.bufPrint(
  97         entry_buffer[0..],
  98         "scan_lookback_{d}x{d}_",
  99         .{ @as(u64, @intCast(total)), expected_blocks },
 100     );
 101     try std.testing.expect(std.mem.indexOf(u8, summary.entry_name, expected_entry_part) != null);
 102 }
 103 
 104 test "numerics: f32 add matches the host reference exactly" {
 105     const allocator = std.testing.allocator;
 106     const module = try binaryF32Module(allocator, .add);
 107 
 108     var lhs = [_]f32{ 1.0, -2.5, 3.25, 0.0, 1e6, -1e-6, 7.5, -8.0 };
 109     var rhs = [_]f32{ 0.5, 2.5, -3.25, -0.0, 1.0, 1e-6, 0.25, 8.0 };
 110     var out = @as([element_count]f32, @splat(0.0));
 111     var count = [_]i32{element_count};
 112 
 113     try runSingleKernel(allocator, module, &.{
 114         kernel.argumentBuffer(f32, out[0..]),
 115         kernel.argumentBuffer(f32, lhs[0..]),
 116         kernel.argumentBuffer(f32, rhs[0..]),
 117         kernel.argumentBuffer(i32, count[0..]),
 118     });
 119 
 120     for (lhs, rhs, out) |a, b, actual| {
 121         try std.testing.expectEqual(a + b, actual);
 122     }
 123 }
 124 
 125 test "numerics: f32 tanh stays within transcendental tolerance" {
 126     const allocator = std.testing.allocator;
 127     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 128     errdefer builder.deinit();
 129     const vec_ty = try builder.tensor(.f32, &.{element_count});
 130     var fb = try builder.beginFunction("numerics_tanh", &.{vec_ty}, &.{vec_ty});
 131     const out_value = try fb.tanh(fb.parameter(0));
 132     try fb.return_(&.{out_value});
 133     try fb.finish();
 134     const module = try builder.finish();
 135 
 136     var input = [_]f32{ -4.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };
 137     var out = @as([element_count]f32, @splat(0.0));
 138     var count = [_]i32{element_count};
 139 
 140     try runSingleKernel(allocator, module, &.{
 141         kernel.argumentBuffer(f32, out[0..]),
 142         kernel.argumentBuffer(f32, input[0..]),
 143         kernel.argumentBuffer(i32, count[0..]),
 144     });
 145 
 146     for (input, out) |x, actual| {
 147         try std.testing.expectApproxEqAbs(std.math.tanh(x), actual, 0.000001);
 148     }
 149 }
 150 
 151 test "numerics: f32 atan2 handles finite quadrants" {
 152     const allocator = std.testing.allocator;
 153     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 154     errdefer builder.deinit();
 155     const vec_ty = try builder.tensor(.f32, &.{element_count});
 156     var fb = try builder.beginFunction("numerics_atan2", &.{ vec_ty, vec_ty }, &.{vec_ty});
 157     const out_value = try fb.atan2(fb.parameter(0), fb.parameter(1));
 158     try fb.return_(&.{out_value});
 159     try fb.finish();
 160     const module = try builder.finish();
 161 
 162     var y = [_]f32{ 1.0, 1.0, -1.0, -1.0, 0.5, -0.5, 2.0, -2.0 };
 163     var x = [_]f32{ 1.0, -1.0, -1.0, 1.0, 2.0, 2.0, -0.5, -0.5 };
 164     var out = @as([element_count]f32, @splat(0.0));
 165     var count = [_]i32{element_count};
 166 
 167     try runSingleKernel(allocator, module, &.{
 168         kernel.argumentBuffer(f32, out[0..]),
 169         kernel.argumentBuffer(f32, y[0..]),
 170         kernel.argumentBuffer(f32, x[0..]),
 171         kernel.argumentBuffer(i32, count[0..]),
 172     });
 173 
 174     for (y, x, out) |yy, xx, actual| {
 175         try std.testing.expectApproxEqAbs(std.math.atan2(yy, xx), actual, 0.000001);
 176     }
 177 }
 178 
 179 test "numerics: f32 round matches the host reference bits" {
 180     const allocator = std.testing.allocator;
 181     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 182     errdefer builder.deinit();
 183     const vec_ty = try builder.tensor(.f32, &.{element_count});
 184     var fb = try builder.beginFunction("numerics_round", &.{vec_ty}, &.{vec_ty});
 185     const out_value = try fb.round(fb.parameter(0));
 186     try fb.return_(&.{out_value});
 187     try fb.finish();
 188     const module = try builder.finish();
 189 
 190     var input = [_]f32{ -1.5, -0.5, -0.25, -0.0, 0.0, 0.25, 0.5, 1.5 };
 191     var out = @as([element_count]f32, @splat(0.0));
 192     var count = [_]i32{element_count};
 193 
 194     try runSingleKernel(allocator, module, &.{
 195         kernel.argumentBuffer(f32, out[0..]),
 196         kernel.argumentBuffer(f32, input[0..]),
 197         kernel.argumentBuffer(i32, count[0..]),
 198     });
 199 
 200     for (input, out) |x, actual| {
 201         try std.testing.expectEqual(@as(u32, @bitCast(@round(x))), @as(u32, @bitCast(actual)));
 202     }
 203 }
 204 
 205 test "numerics: i32 mul matches the host reference exactly" {
 206     const allocator = std.testing.allocator;
 207     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 208     errdefer builder.deinit();
 209     const vec_ty = try builder.tensor(.i32, &.{element_count});
 210     var fb = try builder.beginFunction("numerics_i32_mul", &.{ vec_ty, vec_ty }, &.{vec_ty});
 211     const out_value = try fb.mul(fb.parameter(0), fb.parameter(1));
 212     try fb.return_(&.{out_value});
 213     try fb.finish();
 214     const module = try builder.finish();
 215 
 216     var lhs = [_]i32{ 1, -2, 3, 0, 100000, -7, 11, -13 };
 217     var rhs = [_]i32{ 5, 2, -3, 9, 3, -7, 0, 13 };
 218     var out = @as([element_count]i32, @splat(0));
 219     var count = [_]i32{element_count};
 220 
 221     try runSingleKernel(allocator, module, &.{
 222         kernel.argumentBuffer(i32, out[0..]),
 223         kernel.argumentBuffer(i32, lhs[0..]),
 224         kernel.argumentBuffer(i32, rhs[0..]),
 225         kernel.argumentBuffer(i32, count[0..]),
 226     });
 227 
 228     for (lhs, rhs, out) |a, b, actual| {
 229         try std.testing.expectEqual(a *% b, actual);
 230     }
 231 }
 232 
 233 test "numerics: f32 convert to i32 truncates toward zero" {
 234     const allocator = std.testing.allocator;
 235     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 236     errdefer builder.deinit();
 237     const in_ty = try builder.tensor(.f32, &.{element_count});
 238     const out_ty = try builder.tensor(.i32, &.{element_count});
 239     var fb = try builder.beginFunction("numerics_convert", &.{in_ty}, &.{out_ty});
 240     const out_value = try fb.convert(fb.parameter(0), out_ty, .i32);
 241     try fb.return_(&.{out_value});
 242     try fb.finish();
 243     const module = try builder.finish();
 244 
 245     var input = [_]f32{ 0.0, 1.5, -1.5, 2.9, -2.9, 100.0, -100.0, 0.4 };
 246     var out = @as([element_count]i32, @splat(0));
 247     var count = [_]i32{element_count};
 248 
 249     try runSingleKernel(allocator, module, &.{
 250         kernel.argumentBuffer(i32, out[0..]),
 251         kernel.argumentBuffer(f32, input[0..]),
 252         kernel.argumentBuffer(i32, count[0..]),
 253     });
 254 
 255     for (input, out) |x, actual| {
 256         try std.testing.expectEqual(@as(i32, @intFromFloat(x)), actual);
 257     }
 258 }
 259 
 260 test "numerics: fused f32 convert feeds i32 add" {
 261     const allocator = std.testing.allocator;
 262     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 263     errdefer builder.deinit();
 264     const in_ty = try builder.tensor(.f32, &.{element_count});
 265     const i32_ty = try builder.tensor(.i32, &.{element_count});
 266     var fb = try builder.beginFunction("numerics_fused_convert_i32_add", &.{ in_ty, i32_ty }, &.{i32_ty});
 267     const converted = try fb.convert(fb.parameter(0), i32_ty, .i32);
 268     const out_value = try fb.add(converted, fb.parameter(1));
 269     try fb.return_(&.{out_value});
 270     try fb.finish();
 271     const module = try builder.finish();
 272 
 273     var input = [_]f32{ 0.0, 1.5, -1.5, 2.9, -2.9, 100.0, -100.0, 0.4 };
 274     var bias = [_]i32{ 1, 2, 3, -4, -5, 6, 7, -8 };
 275     var out = @as([element_count]i32, @splat(0));
 276     var count = [_]i32{element_count};
 277 
 278     try runSingleKernel(allocator, module, &.{
 279         kernel.argumentBuffer(i32, out[0..]),
 280         kernel.argumentBuffer(f32, input[0..]),
 281         kernel.argumentBuffer(i32, bias[0..]),
 282         kernel.argumentBuffer(i32, count[0..]),
 283     });
 284 
 285     for (input, bias, out) |x, y, actual| {
 286         try std.testing.expectEqual(@as(i32, @intFromFloat(x)) +% y, actual);
 287     }
 288 }
 289 
 290 test "numerics: fused f16 convert feeds f32 add" {
 291     const allocator = std.testing.allocator;
 292     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 293     errdefer builder.deinit();
 294     const f16_ty = try builder.tensor(.f16, &.{element_count});
 295     const f32_ty = try builder.tensor(.f32, &.{element_count});
 296     var fb = try builder.beginFunction("numerics_fused_convert_f32_add", &.{ f16_ty, f32_ty }, &.{f32_ty});
 297     const converted = try fb.convert(fb.parameter(0), f32_ty, .f32);
 298     const out_value = try fb.add(converted, fb.parameter(1));
 299     try fb.return_(&.{out_value});
 300     try fb.finish();
 301     const module = try builder.finish();
 302 
 303     var input = [_]f16{ 0.5, -1.0, 2.25, 0.0, 1.5, -0.25, 3.0, -2.0 };
 304     var bias = [_]f32{ 0.1, -0.2, 0.3, -0.4, 0.0, 0.5, -0.6, 0.7 };
 305     var out = @as([element_count]f32, @splat(0.0));
 306     var count = [_]i32{element_count};
 307 
 308     try runSingleKernel(allocator, module, &.{
 309         kernel.argumentBuffer(f32, out[0..]),
 310         kernel.argumentBuffer(f16, input[0..]),
 311         kernel.argumentBuffer(f32, bias[0..]),
 312         kernel.argumentBuffer(i32, count[0..]),
 313     });
 314 
 315     for (input, bias, out) |x, y, actual| {
 316         try std.testing.expectApproxEqAbs(@as(f32, x) + y, actual, 0.000001);
 317     }
 318 }
 319 
 320 test "numerics: f32 reduce sum honors non-identity constant init" {
 321     const allocator = std.testing.allocator;
 322     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 323     errdefer builder.deinit();
 324     const in_ty = try builder.tensor(.f32, &.{ 2, 4 });
 325     const init_ty = try builder.tensor(.f32, &.{});
 326     const out_ty = try builder.tensor(.f32, &.{2});
 327     var fb = try builder.beginFunction("numerics_reduce_init", &.{in_ty}, &.{out_ty});
 328     const bias = [_]f32{5.0};
 329     const init = try fb.constant(init_ty, std.mem.sliceAsBytes(bias[0..]));
 330     const out_value = try fb.reduce(fb.parameter(0), init, out_ty, "sum", &.{1});
 331     try fb.return_(&.{out_value});
 332     try fb.finish();
 333     const module = try builder.finish();
 334 
 335     var input = [_]f32{ 1.0, 2.0, 3.0, 4.0, -1.0, -2.0, 0.5, 0.5 };
 336     var out = @as([2]f32, @splat(0.0));
 337     var count = [_]i32{2};
 338 
 339     try runSingleKernel(allocator, module, &.{
 340         kernel.argumentBuffer(f32, out[0..]),
 341         kernel.argumentBuffer(f32, input[0..]),
 342         kernel.argumentBuffer(i32, count[0..]),
 343     });
 344 
 345     try std.testing.expectApproxEqAbs(@as(f32, 15.0), out[0], 0.000001);
 346     try std.testing.expectApproxEqAbs(@as(f32, 3.0), out[1], 0.000001);
 347 }
 348 
 349 test "numerics: f32 reduce sum reads parameter init" {
 350     const allocator = std.testing.allocator;
 351     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 352     errdefer builder.deinit();
 353     const in_ty = try builder.tensor(.f32, &.{ 2, 4 });
 354     const init_ty = try builder.tensor(.f32, &.{});
 355     const out_ty = try builder.tensor(.f32, &.{2});
 356     var fb = try builder.beginFunction("numerics_reduce_parameter_init", &.{ in_ty, init_ty }, &.{out_ty});
 357     const out_value = try fb.reduce(fb.parameter(0), fb.parameter(1), out_ty, "sum", &.{1});
 358     try fb.return_(&.{out_value});
 359     try fb.finish();
 360     const module = try builder.finish();
 361 
 362     var input = [_]f32{ 1.0, 2.0, 3.0, 4.0, -1.0, -2.0, 0.5, 0.5 };
 363     var init = [_]f32{5.0};
 364     var out = @as([2]f32, @splat(0.0));
 365     var count = [_]i32{2};
 366 
 367     try runSingleKernel(allocator, module, &.{
 368         kernel.argumentBuffer(f32, out[0..]),
 369         kernel.argumentBuffer(f32, input[0..]),
 370         kernel.argumentBuffer(f32, init[0..]),
 371         kernel.argumentBuffer(i32, count[0..]),
 372     });
 373 
 374     try std.testing.expectApproxEqAbs(@as(f32, 15.0), out[0], 0.000001);
 375     try std.testing.expectApproxEqAbs(@as(f32, 3.0), out[1], 0.000001);
 376 }
 377 
 378 test "numerics: f16 reduce sum rounds through half precision" {
 379     const allocator = std.testing.allocator;
 380     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 381     errdefer builder.deinit();
 382     const in_ty = try builder.tensor(.f16, &.{ 2, 4 });
 383     const init_ty = try builder.tensor(.f16, &.{});
 384     const out_ty = try builder.tensor(.f16, &.{2});
 385     var fb = try builder.beginFunction("numerics_f16_reduce_sum", &.{ in_ty, init_ty }, &.{out_ty});
 386     const out_value = try fb.reduce(fb.parameter(0), fb.parameter(1), out_ty, "sum", &.{1});
 387     try fb.return_(&.{out_value});
 388     try fb.finish();
 389     const module = try builder.finish();
 390 
 391     var input = [_]f16{ 1.0, 2.0, 3.0, 4.0, -1.0, -2.0, 0.5, 0.5 };
 392     var init = [_]f16{5.0};
 393     var out = @as([2]f16, @splat(0.0));
 394     var count = [_]i32{2};
 395 
 396     try runSingleKernel(allocator, module, &.{
 397         kernel.argumentBuffer(f16, out[0..]),
 398         kernel.argumentBuffer(f16, input[0..]),
 399         kernel.argumentBuffer(f16, init[0..]),
 400         kernel.argumentBuffer(i32, count[0..]),
 401     });
 402 
 403     try std.testing.expectApproxEqAbs(@as(f32, 15.0), @as(f32, out[0]), 0.001);
 404     try std.testing.expectApproxEqAbs(@as(f32, 3.0), @as(f32, out[1]), 0.001);
 405 }
 406 
 407 test "numerics: i32 reduce max reads parameter init" {
 408     const allocator = std.testing.allocator;
 409     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 410     errdefer builder.deinit();
 411     const in_ty = try builder.tensor(.i32, &.{ 2, 4 });
 412     const init_ty = try builder.tensor(.i32, &.{});
 413     const out_ty = try builder.tensor(.i32, &.{2});
 414     var fb = try builder.beginFunction("numerics_i32_reduce_parameter_init", &.{ in_ty, init_ty }, &.{out_ty});
 415     const out_value = try fb.reduce(fb.parameter(0), fb.parameter(1), out_ty, "max", &.{1});
 416     try fb.return_(&.{out_value});
 417     try fb.finish();
 418     const module = try builder.finish();
 419 
 420     var input = [_]i32{ -4, -3, -2, -1, 1, 7, 3, 5 };
 421     var init = [_]i32{0};
 422     var out = @as([2]i32, @splat(0));
 423     var count = [_]i32{2};
 424 
 425     try runSingleKernel(allocator, module, &.{
 426         kernel.argumentBuffer(i32, out[0..]),
 427         kernel.argumentBuffer(i32, input[0..]),
 428         kernel.argumentBuffer(i32, init[0..]),
 429         kernel.argumentBuffer(i32, count[0..]),
 430     });
 431 
 432     try std.testing.expectEqual(@as(i32, 0), out[0]);
 433     try std.testing.expectEqual(@as(i32, 7), out[1]);
 434 }
 435 
 436 test "numerics: f32 fused compare select picks the smaller operand" {
 437     const allocator = std.testing.allocator;
 438     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 439     errdefer builder.deinit();
 440     const vec_ty = try builder.tensor(.f32, &.{element_count});
 441     const pred_ty = try builder.tensor(.i1, &.{element_count});
 442     var fb = try builder.beginFunction("numerics_compare_select", &.{ vec_ty, vec_ty }, &.{vec_ty});
 443     const pred = try fb.compare(fb.parameter(0), fb.parameter(1), pred_ty, .lt);
 444     const out_value = try fb.select(pred, fb.parameter(0), fb.parameter(1));
 445     try fb.return_(&.{out_value});
 446     try fb.finish();
 447     const module = try builder.finish();
 448 
 449     var lhs = [_]f32{ 1.0, -2.5, 3.25, 0.0, 7.0, -8.0, 0.5, 2.0 };
 450     var rhs = [_]f32{ 0.5, 2.5, -3.25, -1.0, 7.5, -9.0, 0.5, 1.0 };
 451     var out = @as([element_count]f32, @splat(0.0));
 452     var count = [_]i32{element_count};
 453 
 454     try runSingleKernel(allocator, module, &.{
 455         kernel.argumentBuffer(f32, out[0..]),
 456         kernel.argumentBuffer(f32, lhs[0..]),
 457         kernel.argumentBuffer(f32, rhs[0..]),
 458         kernel.argumentBuffer(i32, count[0..]),
 459     });
 460 
 461     for (lhs, rhs, out) |a, b, actual| {
 462         try std.testing.expectEqual(@min(a, b), actual);
 463     }
 464 }
 465 
 466 test "numerics: returned compare mask stores byte-sized bools" {
 467     const allocator = std.testing.allocator;
 468     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 469     errdefer builder.deinit();
 470     const vec_ty = try builder.tensor(.f32, &.{element_count});
 471     const pred_ty = try builder.tensor(.i1, &.{element_count});
 472     var fb = try builder.beginFunction("numerics_compare_mask", &.{ vec_ty, vec_ty }, &.{pred_ty});
 473     const out_value = try fb.compare(fb.parameter(0), fb.parameter(1), pred_ty, .lt);
 474     try fb.return_(&.{out_value});
 475     try fb.finish();
 476     const module = try builder.finish();
 477 
 478     var lhs = [_]f32{ 1.0, -2.5, 3.25, 0.0, 7.0, -8.0, 0.5, 2.0 };
 479     var rhs = [_]f32{ 0.5, 2.5, -3.25, -1.0, 7.5, -9.0, 0.5, 1.0 };
 480     var out = @as([element_count]bool, @splat(false));
 481     var count = [_]i32{element_count};
 482 
 483     try runSingleKernel(allocator, module, &.{
 484         kernel.argumentBuffer(bool, out[0..]),
 485         kernel.argumentBuffer(f32, lhs[0..]),
 486         kernel.argumentBuffer(f32, rhs[0..]),
 487         kernel.argumentBuffer(i32, count[0..]),
 488     });
 489 
 490     for (lhs, rhs, out) |a, b, actual| {
 491         try std.testing.expectEqual(a < b, actual);
 492     }
 493 }
 494 
 495 test "numerics: external bool mask feeds select" {
 496     const allocator = std.testing.allocator;
 497     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 498     errdefer builder.deinit();
 499     const vec_ty = try builder.tensor(.f32, &.{element_count});
 500     const pred_ty = try builder.tensor(.i1, &.{element_count});
 501     var fb = try builder.beginFunction("numerics_external_mask_select", &.{ pred_ty, vec_ty, vec_ty }, &.{vec_ty});
 502     const out_value = try fb.select(fb.parameter(0), fb.parameter(1), fb.parameter(2));
 503     try fb.return_(&.{out_value});
 504     try fb.finish();
 505     const module = try builder.finish();
 506 
 507     var mask = [_]bool{ true, false, true, false, false, true, true, false };
 508     var on_true = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 };
 509     var on_false = [_]f32{ -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0 };
 510     var out = @as([element_count]f32, @splat(0.0));
 511     var count = [_]i32{element_count};
 512 
 513     try runSingleKernel(allocator, module, &.{
 514         kernel.argumentBuffer(f32, out[0..]),
 515         kernel.argumentBuffer(bool, mask[0..]),
 516         kernel.argumentBuffer(f32, on_true[0..]),
 517         kernel.argumentBuffer(f32, on_false[0..]),
 518         kernel.argumentBuffer(i32, count[0..]),
 519     });
 520 
 521     for (mask, on_true, on_false, out) |take_true, yes, no, actual| {
 522         try std.testing.expectEqual(if (take_true) yes else no, actual);
 523     }
 524 }
 525 
 526 test "numerics: ordered compare treats NaN as false and select takes the other branch" {
 527     const allocator = std.testing.allocator;
 528     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 529     errdefer builder.deinit();
 530     const vec_ty = try builder.tensor(.f32, &.{element_count});
 531     const pred_ty = try builder.tensor(.i1, &.{element_count});
 532     var fb = try builder.beginFunction("numerics_nan_compare", &.{ vec_ty, vec_ty }, &.{vec_ty});
 533     const pred = try fb.compare(fb.parameter(0), fb.parameter(1), pred_ty, .lt);
 534     const out_value = try fb.select(pred, fb.parameter(0), fb.parameter(1));
 535     try fb.return_(&.{out_value});
 536     try fb.finish();
 537     const module = try builder.finish();
 538 
 539     const nan = std.math.nan(f32);
 540     const inf = std.math.inf(f32);
 541     var lhs = [_]f32{ nan, 1.0, nan, -inf, inf, 0.0, nan, 2.0 };
 542     var rhs = [_]f32{ 1.0, nan, nan, 0.0, inf, -0.0, -inf, 3.0 };
 543     var out = @as([element_count]f32, @splat(0.0));
 544     var count = [_]i32{element_count};
 545 
 546     try runSingleKernel(allocator, module, &.{
 547         kernel.argumentBuffer(f32, out[0..]),
 548         kernel.argumentBuffer(f32, lhs[0..]),
 549         kernel.argumentBuffer(f32, rhs[0..]),
 550         kernel.argumentBuffer(i32, count[0..]),
 551     });
 552 
 553     for (lhs, rhs, out) |a, b, actual| {
 554         const expected = if (a < b) a else b;
 555         if (std.math.isNan(expected)) {
 556             try std.testing.expect(std.math.isNan(actual));
 557         } else {
 558             try std.testing.expectEqual(expected, actual);
 559         }
 560     }
 561 }
 562 
 563 test "numerics: f32 min and max propagate NaN like the host" {
 564     const allocator = std.testing.allocator;
 565     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 566     errdefer builder.deinit();
 567     const vec_ty = try builder.tensor(.f32, &.{element_count});
 568     var fb = try builder.beginFunction("numerics_minmax_nan", &.{ vec_ty, vec_ty }, &.{vec_ty});
 569     const low = try fb.min(fb.parameter(0), fb.parameter(1));
 570     const out_value = try fb.max(low, fb.parameter(1));
 571     try fb.return_(&.{out_value});
 572     try fb.finish();
 573     const module = try builder.finish();
 574 
 575     const nan = std.math.nan(f32);
 576     const inf = std.math.inf(f32);
 577     var lhs = [_]f32{ nan, 1.0, -inf, inf, 0.0, -1.0, nan, 4.0 };
 578     var rhs = [_]f32{ 1.0, nan, 0.0, inf, -0.0, -2.0, nan, 3.0 };
 579     var out = @as([element_count]f32, @splat(0.0));
 580     var count = [_]i32{element_count};
 581 
 582     try runSingleKernel(allocator, module, &.{
 583         kernel.argumentBuffer(f32, out[0..]),
 584         kernel.argumentBuffer(f32, lhs[0..]),
 585         kernel.argumentBuffer(f32, rhs[0..]),
 586         kernel.argumentBuffer(i32, count[0..]),
 587     });
 588 
 589     for (lhs, rhs, out) |a, b, actual| {
 590         const expected = @max(@min(a, b), b);
 591         if (std.math.isNan(expected)) {
 592             try std.testing.expect(std.math.isNan(actual));
 593         } else {
 594             try std.testing.expectEqual(expected, actual);
 595         }
 596     }
 597 }
 598 
 599 test "numerics: transpose moves every element exactly" {
 600     const allocator = std.testing.allocator;
 601     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 602     errdefer builder.deinit();
 603     const in_ty = try builder.tensor(.f32, &.{ 2, 4 });
 604     const out_ty = try builder.tensor(.f32, &.{ 4, 2 });
 605     var fb = try builder.beginFunction("numerics_transpose", &.{in_ty}, &.{out_ty});
 606     const out_value = try fb.transpose(fb.parameter(0), out_ty, &.{ 1, 0 });
 607     try fb.return_(&.{out_value});
 608     try fb.finish();
 609     const module = try builder.finish();
 610 
 611     var input = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 };
 612     var out = @as([element_count]f32, @splat(0.0));
 613     var count = [_]i32{element_count};
 614 
 615     try runSingleKernel(allocator, module, &.{
 616         kernel.argumentBuffer(f32, out[0..]),
 617         kernel.argumentBuffer(f32, input[0..]),
 618         kernel.argumentBuffer(i32, count[0..]),
 619     });
 620 
 621     try std.testing.expectEqualSlices(f32, &.{ 1.0, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0 }, out[0..]);
 622 }
 623 
 624 test "numerics: reshape preserves row-major element order exactly" {
 625     const allocator = std.testing.allocator;
 626     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 627     errdefer builder.deinit();
 628     const in_ty = try builder.tensor(.f32, &.{ 2, 4 });
 629     const out_ty = try builder.tensor(.f32, &.{ 4, 2 });
 630     var fb = try builder.beginFunction("numerics_reshape", &.{in_ty}, &.{out_ty});
 631     const out_value = try fb.reshape(fb.parameter(0), out_ty, &.{ 4, 2 });
 632     try fb.return_(&.{out_value});
 633     try fb.finish();
 634     const module = try builder.finish();
 635 
 636     var input = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 };
 637     var out = @as([element_count]f32, @splat(0.0));
 638     var count = [_]i32{element_count};
 639 
 640     try runSingleKernel(allocator, module, &.{
 641         kernel.argumentBuffer(f32, out[0..]),
 642         kernel.argumentBuffer(f32, input[0..]),
 643         kernel.argumentBuffer(i32, count[0..]),
 644     });
 645 
 646     try std.testing.expectEqualSlices(f32, input[0..], out[0..]);
 647 }
 648 
 649 test "numerics: slice strides select exact elements" {
 650     const allocator = std.testing.allocator;
 651     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 652     errdefer builder.deinit();
 653     const in_ty = try builder.tensor(.f32, &.{element_count});
 654     const out_ty = try builder.tensor(.f32, &.{4});
 655     var fb = try builder.beginFunction("numerics_slice", &.{in_ty}, &.{out_ty});
 656     const out_value = try fb.slice(fb.parameter(0), out_ty, &.{1}, &.{8}, &.{2});
 657     try fb.return_(&.{out_value});
 658     try fb.finish();
 659     const module = try builder.finish();
 660 
 661     var input = [_]f32{ 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0 };
 662     var out = @as([4]f32, @splat(0.0));
 663     var count = [_]i32{4};
 664 
 665     try runSingleKernel(allocator, module, &.{
 666         kernel.argumentBuffer(f32, out[0..]),
 667         kernel.argumentBuffer(f32, input[0..]),
 668         kernel.argumentBuffer(i32, count[0..]),
 669     });
 670 
 671     try std.testing.expectEqualSlices(f32, &.{ 11.0, 13.0, 15.0, 17.0 }, out[0..]);
 672 }
 673 
 674 test "numerics: gather selects exactly the indexed elements" {
 675     const allocator = std.testing.allocator;
 676     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 677     errdefer builder.deinit();
 678     const in_ty = try builder.tensor(.f32, &.{element_count});
 679     const indices_ty = try builder.tensor(.i32, &.{4});
 680     const out_ty = try builder.tensor(.f32, &.{4});
 681     var fb = try builder.beginFunction("numerics_gather", &.{ in_ty, indices_ty }, &.{out_ty});
 682     const out_value = try fb.gather(fb.parameter(0), fb.parameter(1), out_ty, 0);
 683     try fb.return_(&.{out_value});
 684     try fb.finish();
 685     const module = try builder.finish();
 686 
 687     var input = [_]f32{ 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0 };
 688     var indices = [_]i32{ 7, 0, 3, 3 };
 689     var out = @as([4]f32, @splat(0.0));
 690     var count = [_]i32{4};
 691 
 692     try runSingleKernel(allocator, module, &.{
 693         kernel.argumentBuffer(f32, out[0..]),
 694         kernel.argumentBuffer(f32, input[0..]),
 695         kernel.argumentBuffer(i32, indices[0..]),
 696         kernel.argumentBuffer(i32, count[0..]),
 697     });
 698 
 699     try std.testing.expectEqualSlices(f32, &.{ 17.0, 10.0, 13.0, 13.0 }, out[0..]);
 700 }
 701 
 702 test "numerics: pad fills low and high edges exactly" {
 703     const allocator = std.testing.allocator;
 704     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 705     errdefer builder.deinit();
 706     const in_ty = try builder.tensor(.f32, &.{4});
 707     const fill_ty = try builder.tensor(.f32, &.{});
 708     const out_ty = try builder.tensor(.f32, &.{8});
 709     var fb = try builder.beginFunction("numerics_pad", &.{ in_ty, fill_ty }, &.{out_ty});
 710     const out_value = try fb.pad(fb.parameter(0), fb.parameter(1), out_ty, &.{2}, &.{2}, &.{0});
 711     try fb.return_(&.{out_value});
 712     try fb.finish();
 713     const module = try builder.finish();
 714 
 715     var input = [_]f32{ 1.0, 2.0, 3.0, 4.0 };
 716     var fill = [_]f32{-1.0};
 717     var out = @as([element_count]f32, @splat(0.0));
 718     var count = [_]i32{element_count};
 719 
 720     try runSingleKernel(allocator, module, &.{
 721         kernel.argumentBuffer(f32, out[0..]),
 722         kernel.argumentBuffer(f32, input[0..]),
 723         kernel.argumentBuffer(f32, fill[0..]),
 724         kernel.argumentBuffer(i32, count[0..]),
 725     });
 726 
 727     try std.testing.expectEqualSlices(f32, &.{ -1.0, -1.0, 1.0, 2.0, 3.0, 4.0, -1.0, -1.0 }, out[0..]);
 728 }
 729 
 730 test "numerics: concatenate appends operands exactly" {
 731     const allocator = std.testing.allocator;
 732     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 733     errdefer builder.deinit();
 734     const vec_ty = try builder.tensor(.f32, &.{4});
 735     const out_ty = try builder.tensor(.f32, &.{8});
 736     var fb = try builder.beginFunction("numerics_concatenate", &.{ vec_ty, vec_ty }, &.{out_ty});
 737     const out_value = try fb.concatenate(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, 0);
 738     try fb.return_(&.{out_value});
 739     try fb.finish();
 740     const module = try builder.finish();
 741 
 742     var lhs = [_]f32{ 1.0, 2.0, 3.0, 4.0 };
 743     var rhs = [_]f32{ 5.0, 6.0, 7.0, 8.0 };
 744     var out = @as([element_count]f32, @splat(0.0));
 745     var count = [_]i32{element_count};
 746 
 747     try runSingleKernel(allocator, module, &.{
 748         kernel.argumentBuffer(f32, out[0..]),
 749         kernel.argumentBuffer(f32, lhs[0..]),
 750         kernel.argumentBuffer(f32, rhs[0..]),
 751         kernel.argumentBuffer(i32, count[0..]),
 752     });
 753 
 754     try std.testing.expectEqualSlices(f32, &.{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 }, out[0..]);
 755 }
 756 
 757 test "numerics: scatter writes updates at exact indices" {
 758     const allocator = std.testing.allocator;
 759     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 760     errdefer builder.deinit();
 761     const in_ty = try builder.tensor(.f32, &.{8});
 762     const indices_ty = try builder.tensor(.i32, &.{4});
 763     const updates_ty = try builder.tensor(.f32, &.{4});
 764     var fb = try builder.beginFunction("numerics_scatter", &.{ in_ty, indices_ty, updates_ty }, &.{in_ty});
 765     const out_value = try fb.scatter(fb.parameter(0), fb.parameter(1), fb.parameter(2), in_ty, 0);
 766     try fb.return_(&.{out_value});
 767     try fb.finish();
 768     const module = try builder.finish();
 769 
 770     var input = [_]f32{ 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0 };
 771     var indices = [_]i32{ 6, 1, 4, 3 };
 772     var updates = [_]f32{ 60.0, 10.0, 40.0, 30.0 };
 773     var out = @as([element_count]f32, @splat(0.0));
 774     var count = [_]i32{element_count};
 775 
 776     try runSingleKernel(allocator, module, &.{
 777         kernel.argumentBuffer(f32, out[0..]),
 778         kernel.argumentBuffer(f32, input[0..]),
 779         kernel.argumentBuffer(i32, indices[0..]),
 780         kernel.argumentBuffer(f32, updates[0..]),
 781         kernel.argumentBuffer(i32, count[0..]),
 782     });
 783 
 784     try std.testing.expectEqualSlices(f32, &.{ 10.0, 10.0, 12.0, 30.0, 40.0, 15.0, 60.0, 17.0 }, out[0..]);
 785 }
 786 
 787 test "numerics: fused elementwise chain lowers to one kernel and computes exactly" {
 788     const allocator = std.testing.allocator;
 789     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 790     errdefer builder.deinit();
 791     const vec_ty = try builder.tensor(.f32, &.{element_count});
 792     var fb = try builder.beginFunction("numerics_fused_chain", &.{ vec_ty, vec_ty, vec_ty }, &.{vec_ty});
 793     const product = try fb.mul(fb.parameter(0), fb.parameter(1));
 794     const shifted = try fb.add(product, fb.parameter(2));
 795     const out_value = try fb.tanh(shifted);
 796     try fb.return_(&.{out_value});
 797     try fb.finish();
 798     const module = try builder.finish();
 799 
 800     var a = [_]f32{ 0.5, -1.0, 2.0, 0.0, 1.5, -0.25, 3.0, -2.0 };
 801     var b = [_]f32{ 2.0, 0.5, -1.0, 4.0, 1.0, 8.0, -0.5, 0.75 };
 802     var c = [_]f32{ 0.1, -0.2, 0.3, -0.4, 0.0, 0.5, -0.6, 0.7 };
 803     var out = @as([element_count]f32, @splat(0.0));
 804     var count = [_]i32{element_count};
 805 
 806     try runSingleKernel(allocator, module, &.{
 807         kernel.argumentBuffer(f32, out[0..]),
 808         kernel.argumentBuffer(f32, a[0..]),
 809         kernel.argumentBuffer(f32, b[0..]),
 810         kernel.argumentBuffer(f32, c[0..]),
 811         kernel.argumentBuffer(i32, count[0..]),
 812     });
 813 
 814     for (a, b, c, out) |x, y, z, actual| {
 815         try std.testing.expectApproxEqAbs(std.math.tanh(x * y + z), actual, 0.000001);
 816     }
 817 }
 818 
 819 test "numerics: iterate runs per-element do-while trip counts" {
 820     const allocator = std.testing.allocator;
 821 
 822     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 823     errdefer builder.deinit();
 824     const vec_ty = try builder.tensor(.f32, &.{element_count});
 825     const pred_ty = try builder.tensor(.i1, &.{element_count});
 826     var fb = try builder.beginFunction("numerics_iterate", &.{vec_ty}, &.{vec_ty});
 827 
 828     var it = try fb.beginIterate(&.{fb.parameter(0)}, 10);
 829     const body = it.inner();
 830     var ones: [element_count]f32 = @splat(1.0);
 831     var limits: [element_count]f32 = @splat(5.0);
 832     const one = try body.constant(vec_ty, std.mem.sliceAsBytes(ones[0..]));
 833     const limit = try body.constant(vec_ty, std.mem.sliceAsBytes(limits[0..]));
 834     const stepped = try body.add(it.carry(0), one);
 835     const active = try body.compare(stepped, limit, pred_ty, .lt);
 836     try it.yield_(active, &.{stepped});
 837 
 838     try fb.return_(&.{it.result(0)});
 839     try fb.finish();
 840     const module = try builder.finish();
 841 
 842     var input = [_]f32{ 0.0, 3.5, 4.0, 4.5, 7.0, -100.0, 4.999, -2.25 };
 843     var out = @as([element_count]f32, @splat(0.0));
 844     var count = [_]i32{element_count};
 845 
 846     try runSingleKernel(allocator, module, &.{
 847         kernel.argumentBuffer(f32, out[0..]),
 848         kernel.argumentBuffer(f32, input[0..]),
 849         kernel.argumentBuffer(i32, count[0..]),
 850     });
 851 
 852     for (input, out) |start, actual| {
 853         var value = start;
 854         var iters: usize = 0;
 855         while (iters < 10) : (iters += 1) {
 856             value += 1.0;
 857             if (!(value < 5.0)) break;
 858         }
 859         try std.testing.expectApproxEqAbs(value, actual, 0.00001);
 860     }
 861 }
 862 
 863 test "numerics: cumsum lookback kernel matches the host on CUDA" {
 864     const allocator = std.testing.allocator;
 865     var state = gpu.cuda.State.initDevice(allocator, 0) catch return error.SkipZigTest;
 866     defer state.deinit();
 867 
 868     const total = 16384;
 869     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 870     errdefer builder.deinit();
 871     const vec_ty = try builder.tensor(.f32, &.{total});
 872     var fb = try builder.beginFunction("probe_scan", &.{vec_ty}, &.{vec_ty});
 873     const out_value = try fb.cumsum(fb.parameter(0), vec_ty, 0);
 874     try fb.return_(&.{out_value});
 875     try fb.finish();
 876     const module = try builder.finish();
 877 
 878     var fragment = try compileAndLoadSemanticModule(
 879         allocator,
 880         state.handle(),
 881         module,
 882         .{},
 883     );
 884     defer fragment.deinit();
 885     try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
 886 
 887     var input: [total]f32 = undefined;
 888     for (input[0..], 0..) |*value, index| value.* = @sin(@as(f32, @floatFromInt(index)) * 0.01) + 0.5;
 889     var out = @as([total]f32, @splat(0.0));
 890     var outputs = [_][]u8{std.mem.sliceAsBytes(out[0..])};
 891     try executable.invoke(fragment, allocator, allocator, &.{std.mem.sliceAsBytes(input[0..])}, &outputs);
 892 
 893     var acc: f64 = 0;
 894     var worst: f64 = 0;
 895     for (0..total) |i| {
 896         acc += input[i];
 897         worst = @max(worst, @abs(acc - out[i]));
 898     }
 899     try std.testing.expect(worst < 0.01);
 900 }
 901 
 902 test "numerics: cumsum generated scan chooses adaptive block schedule" {
 903     try expectCumsumGeneratedSchedule(4096, 256, 1);
 904     try expectCumsumGeneratedSchedule(8192, 512, 1);
 905     try expectCumsumGeneratedSchedule(12288, 256, 3);
 906 }
 907 
 908 test "numerics: cumsum scans a vector to one kernel and matches the host" {
 909     const allocator = std.testing.allocator;
 910     const total = 1000;
 911 
 912     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 913     errdefer builder.deinit();
 914     const vec_ty = try builder.tensor(.f32, &.{total});
 915     var fb = try builder.beginFunction("numerics_cumsum", &.{vec_ty}, &.{vec_ty});
 916     const out_value = try fb.cumsum(fb.parameter(0), vec_ty, 0);
 917     try fb.return_(&.{out_value});
 918     try fb.finish();
 919     const module = try builder.finish();
 920 
 921     var input: [total]f32 = undefined;
 922     for (input[0..], 0..) |*value, index| {
 923         value.* = @sin(@as(f32, @floatFromInt(index)) * 0.1) + 0.25;
 924     }
 925     var out = @as([total]f32, @splat(0.0));
 926     var count = [_]i32{total};
 927 
 928     try runSingleKernel(allocator, module, &.{
 929         kernel.argumentBuffer(f32, out[0..]),
 930         kernel.argumentBuffer(f32, input[0..]),
 931         kernel.argumentBuffer(i32, count[0..]),
 932     });
 933 
 934     var acc: f64 = 0;
 935     for (0..total) |i| {
 936         acc += input[i];
 937         try std.testing.expectApproxEqAbs(@as(f32, @floatCast(acc)), out[i], 0.002);
 938     }
 939 }
 940 
 941 const flash_attention_seq = 128;
 942 const flash_attention_head = 64;
 943 
 944 fn flashAttentionModule(allocator: std.mem.Allocator) !*SemanticModule {
 945     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
 946     errdefer builder.deinit();
 947     const qkv_ty = try builder.tensor(.f32, &.{ flash_attention_seq, flash_attention_head });
 948     const kt_ty = try builder.tensor(.f32, &.{ flash_attention_head, flash_attention_seq });
 949     const scores_ty = try builder.tensor(.f32, &.{ flash_attention_seq, flash_attention_seq });
 950     const row_ty = try builder.tensor(.f32, &.{flash_attention_seq});
 951     const scalar_ty = try builder.tensor(.f32, &.{});
 952     var fb = try builder.beginFunction("numerics_flash", &.{ qkv_ty, kt_ty, qkv_ty }, &.{qkv_ty});
 953     const q = fb.parameter(0);
 954     const kt = fb.parameter(1);
 955     const v = fb.parameter(2);
 956     const raw_scores = try fb.dotGeneral(q, kt, scores_ty, &.{1}, &.{0}, &.{}, &.{});
 957     const scale_scalar = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.125)));
 958     const scale = try fb.broadcast(scale_scalar, scores_ty, &.{ flash_attention_seq, flash_attention_seq });
 959     const scores = try fb.mul(raw_scores, scale);
 960     const lowest = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, -std.math.floatMax(f32))));
 961     const row_max = try fb.reduce(scores, lowest, row_ty, "max", &.{1});
 962     const shifted = try fb.sub(scores, try fb.broadcastInDim(row_max, scores_ty, &.{ flash_attention_seq, flash_attention_seq }, &.{0}));
 963     const exps = try fb.exp(shifted);
 964     const zero = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
 965     const row_sum = try fb.reduce(exps, zero, row_ty, "sum", &.{1});
 966     const ones_data = @as([flash_attention_seq]f32, @splat(1.0));
 967     const ones = try fb.constant(row_ty, std.mem.sliceAsBytes(ones_data[0..]));
 968     const inverse_sum = try fb.div(ones, row_sum);
 969     const probs = try fb.mul(exps, try fb.broadcastInDim(inverse_sum, scores_ty, &.{ flash_attention_seq, flash_attention_seq }, &.{0}));
 970     const out_value = try fb.dotGeneral(probs, v, qkv_ty, &.{1}, &.{0}, &.{}, &.{});
 971     try fb.return_(&.{out_value});
 972     try fb.finish();
 973     return try builder.finish();
 974 }
 975 
 976 test "numerics: reciprocal-normalized flash attention has one CUDA schedule" {
 977     const allocator = std.testing.allocator;
 978 
 979     var recording = gpu.recording.BackendState{
 980         .allocator = allocator,
 981         .kind = .cuda,
 982         .format = .cuda_ptx,
 983     };
 984     const module = try flashAttentionModule(allocator);
 985     var prepared = try executable.prepareFragmentFromSemanticModule(
 986         allocator,
 987         recording.handle(),
 988         module,
 989         .{},
 990     );
 991     defer prepared.deinit();
 992     try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
 993     const generated = try prepared.generatedKernelSummary(0);
 994     try std.testing.expectEqual(kernelization.GeneratedScheduleKind.matrix, generated.schedule.kind);
 995     try std.testing.expectEqual(@as(u32, 16), generated.schedule.threads.x);
 996     try std.testing.expectEqual(@as(u32, 8), generated.schedule.threads.y);
 997     try std.testing.expect(std.mem.indexOf(u8, generated.entry_name, "flash_attention") != null);
 998 }
 999 
1000 test "numerics: iterate carries multiple states with per-element exits" {
1001     const allocator = std.testing.allocator;
1002 
1003     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1004     errdefer builder.deinit();
1005     const vec_ty = try builder.tensor(.f32, &.{element_count});
1006     const pred_ty = try builder.tensor(.i1, &.{element_count});
1007     var fb = try builder.beginFunction("numerics_iterate_pair", &.{ vec_ty, vec_ty }, &.{ vec_ty, vec_ty });
1008 
1009     var it = try fb.beginIterate(&.{ fb.parameter(0), fb.parameter(1) }, 12);
1010     const body = it.inner();
1011     var halves: [element_count]f32 = @splat(0.5);
1012     var limits: [element_count]f32 = @splat(20.0);
1013     const half = try body.constant(vec_ty, std.mem.sliceAsBytes(halves[0..]));
1014     const limit = try body.constant(vec_ty, std.mem.sliceAsBytes(limits[0..]));
1015     const stepped = try body.add(it.carry(0), it.carry(1));
1016     const decayed = try body.mul(it.carry(1), half);
1017     const active = try body.compare(stepped, limit, pred_ty, .lt);
1018     try it.yield_(active, &.{ stepped, decayed });
1019 
1020     try fb.return_(&.{ it.result(0), it.result(1) });
1021     try fb.finish();
1022     const module = try builder.finish();
1023 
1024     var a = [_]f32{ 0.0, 5.0, 18.0, 30.0, 1.0, -4.0, 10.0, 19.0 };
1025     var b = [_]f32{ 4.0, 8.0, 1.0, 2.0, 16.0, 3.0, 0.5, 2.0 };
1026     var out_a = @as([element_count]f32, @splat(0.0));
1027     var out_b = @as([element_count]f32, @splat(0.0));
1028     var count = [_]i32{element_count};
1029 
1030     const contract_module = try pipeline.prepareContractJobFromSemanticModule(allocator, module, .{});
1031     const tensor_module = try pipeline.prepareTensorJobFromContractJob(allocator, contract_module, .{});
1032     const dispatch_module = try pipeline.prepareDispatchJobFromTensorJob(allocator, tensor_module, .{});
1033     const memory_module = try pipeline.prepareMemoryJobFromDispatchJob(allocator, dispatch_module, .{});
1034     defer memory_module.deinit();
1035 
1036     var pass_ctx = memory_module.passContext();
1037     defer pass_ctx.deinit();
1038     const analysis = try kernelization.getKernelizationAnalysis(&pass_ctx, memory_module.choir_module);
1039     try std.testing.expectEqual(@as(usize, 1), analysis.kernelCount());
1040     try runGeneratedProgram(allocator, &analysis.kernels.items[0].program, &.{
1041         kernel.argumentBuffer(f32, out_a[0..]),
1042         kernel.argumentBuffer(f32, a[0..]),
1043         kernel.argumentBuffer(f32, b[0..]),
1044         kernel.argumentBuffer(f32, out_b[0..]),
1045         kernel.argumentBuffer(i32, count[0..]),
1046     });
1047 
1048     for (0..element_count) |i| {
1049         var av = a[i];
1050         var bv = b[i];
1051         var iters: usize = 0;
1052         while (iters < 12) : (iters += 1) {
1053             const na = av + bv;
1054             const nb = bv * 0.5;
1055             av = na;
1056             bv = nb;
1057             if (!(av < 20.0)) break;
1058         }
1059         try std.testing.expectApproxEqAbs(av, out_a[i], 0.0001);
1060         try std.testing.expectApproxEqAbs(bv, out_b[i], 0.0001);
1061     }
1062 }
1063 
1064 test "numerics: layernorm row pipeline lowers to one kernel and matches the host" {
1065     const allocator = std.testing.allocator;
1066     const rows = 2;
1067     const cols = 1024;
1068 
1069     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1070     errdefer builder.deinit();
1071     const mat_ty = try builder.tensor(.f32, &.{ rows, cols });
1072     const row_ty = try builder.tensor(.f32, &.{rows});
1073     const col_ty = try builder.tensor(.f32, &.{cols});
1074     const scalar_ty = try builder.tensor(.f32, &.{});
1075     var fb = try builder.beginFunction("numerics_layernorm_rows", &.{ mat_ty, col_ty, col_ty }, &.{mat_ty});
1076     const x = fb.parameter(0);
1077     const gamma = fb.parameter(1);
1078     const beta = fb.parameter(2);
1079 
1080     const zero = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
1081     const row_sum = try fb.reduce(x, zero, row_ty, "sum", &.{1});
1082     const inv_cols_values = @as([rows]f32, @splat(1.0 / @as(f32, cols)));
1083     const inv_cols = try fb.constant(row_ty, std.mem.sliceAsBytes(inv_cols_values[0..]));
1084     const mean = try fb.mul(row_sum, inv_cols);
1085     const mean_b = try fb.broadcastInDim(mean, mat_ty, &.{ rows, cols }, &.{0});
1086     const centered = try fb.sub(x, mean_b);
1087     const centered_squared = try fb.mul(centered, centered);
1088     const square_sum = try fb.reduce(centered_squared, zero, row_ty, "sum", &.{1});
1089     const variance = try fb.mul(square_sum, inv_cols);
1090     const eps_values = @as([rows]f32, @splat(0.00001));
1091     const eps = try fb.constant(row_ty, std.mem.sliceAsBytes(eps_values[0..]));
1092     const variance_eps = try fb.add(variance, eps);
1093     const stddev = try fb.sqrt(variance_eps);
1094     const one_values = @as([rows]f32, @splat(1.0));
1095     const one = try fb.constant(row_ty, std.mem.sliceAsBytes(one_values[0..]));
1096     const inv_std = try fb.div(one, stddev);
1097     const inv_std_b = try fb.broadcastInDim(inv_std, mat_ty, &.{ rows, cols }, &.{0});
1098     const normalized = try fb.mul(centered, inv_std_b);
1099     const gamma_b = try fb.broadcastInDim(gamma, mat_ty, &.{ rows, cols }, &.{1});
1100     const scaled = try fb.mul(normalized, gamma_b);
1101     const beta_b = try fb.broadcastInDim(beta, mat_ty, &.{ rows, cols }, &.{1});
1102     const out_value = try fb.add(scaled, beta_b);
1103     try fb.return_(&.{out_value});
1104     try fb.finish();
1105     const module = try builder.finish();
1106 
1107     var input: [rows * cols]f32 = undefined;
1108     for (input[0..], 0..) |*value, index| {
1109         const col: f32 = @floatFromInt(index % cols);
1110         const row: f32 = @floatFromInt(index / cols);
1111         value.* = @sin(col * 0.02) * 1.5 + row * 0.25 + col * 0.001;
1112     }
1113     var gamma_values: [cols]f32 = undefined;
1114     var beta_values: [cols]f32 = undefined;
1115     for (0..cols) |col| {
1116         gamma_values[col] = 0.75 + 0.5 * @as(f32, @floatFromInt(col % 7)) / 7.0;
1117         beta_values[col] = -0.25 + 0.1 * @as(f32, @floatFromInt(col % 5));
1118     }
1119     var out = @as([(rows * cols)]f32, @splat(0.0));
1120     var count = [_]i32{rows * cols};
1121 
1122     try runSingleKernel(allocator, module, &.{
1123         kernel.argumentBuffer(f32, out[0..]),
1124         kernel.argumentBuffer(f32, input[0..]),
1125         kernel.argumentBuffer(f32, gamma_values[0..]),
1126         kernel.argumentBuffer(f32, beta_values[0..]),
1127         kernel.argumentBuffer(i32, count[0..]),
1128     });
1129 
1130     for (0..rows) |row| {
1131         var mean_acc: f32 = 0;
1132         for (0..cols) |col| mean_acc += input[row * cols + col];
1133         const mean_host = mean_acc / cols;
1134         var var_acc: f32 = 0;
1135         for (0..cols) |col| {
1136             const centered_host = input[row * cols + col] - mean_host;
1137             var_acc += centered_host * centered_host;
1138         }
1139         const inv_std_host = 1.0 / @sqrt(var_acc / cols + 0.00001);
1140         for (0..cols) |col| {
1141             const expected = (input[row * cols + col] - mean_host) * inv_std_host * gamma_values[col] + beta_values[col];
1142             try std.testing.expectApproxEqAbs(expected, out[row * cols + col], 0.001);
1143         }
1144     }
1145 }
1146 
1147 test "numerics: softmax row pipeline lowers to one kernel and matches the host" {
1148     const allocator = std.testing.allocator;
1149     const rows = 2;
1150     const cols = 1024;
1151 
1152     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1153     errdefer builder.deinit();
1154     const mat_ty = try builder.tensor(.f32, &.{ rows, cols });
1155     const row_ty = try builder.tensor(.f32, &.{rows});
1156     const scalar_ty = try builder.tensor(.f32, &.{});
1157     var fb = try builder.beginFunction("numerics_softmax_rows", &.{mat_ty}, &.{mat_ty});
1158     const neg_inf = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, -std.math.inf(f32))));
1159     const zero = try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
1160     const row_max = try fb.reduce(fb.parameter(0), neg_inf, row_ty, "max", &.{1});
1161     const max_full = try fb.broadcastInDim(row_max, mat_ty, &.{ rows, cols }, &.{0});
1162     const shifted = try fb.sub(fb.parameter(0), max_full);
1163     const exps = try fb.exp(shifted);
1164     const row_sum = try fb.reduce(exps, zero, row_ty, "sum", &.{1});
1165     const sum_full = try fb.broadcastInDim(row_sum, mat_ty, &.{ rows, cols }, &.{0});
1166     const out_value = try fb.div(exps, sum_full);
1167     try fb.return_(&.{out_value});
1168     try fb.finish();
1169     const module = try builder.finish();
1170 
1171     var input: [rows * cols]f32 = undefined;
1172     for (input[0..], 0..) |*value, index| {
1173         const row: f32 = @floatFromInt(index / cols);
1174         const col: f32 = @floatFromInt(index % cols);
1175         value.* = @sin(col * 0.01) * 2.0 + row * 0.5 - col * 0.0005;
1176     }
1177     var out = @as([(rows * cols)]f32, @splat(0.0));
1178     var count = [_]i32{rows * cols};
1179 
1180     try runSingleKernel(allocator, module, &.{
1181         kernel.argumentBuffer(f32, out[0..]),
1182         kernel.argumentBuffer(f32, input[0..]),
1183         kernel.argumentBuffer(i32, count[0..]),
1184     });
1185 
1186     for (0..rows) |row| {
1187         var host_max: f32 = -std.math.inf(f32);
1188         for (0..cols) |col| {
1189             host_max = @max(host_max, input[row * cols + col]);
1190         }
1191         var host_sum: f32 = 0.0;
1192         for (0..cols) |col| {
1193             host_sum += @exp(input[row * cols + col] - host_max);
1194         }
1195         for (0..cols) |col| {
1196             const expected = @exp(input[row * cols + col] - host_max) / host_sum;
1197             try std.testing.expectApproxEqAbs(expected, out[row * cols + col], 0.00001);
1198         }
1199     }
1200 }
1201 
1202 fn gradLoss(_: *tensor.trace.Builder, args: []const tensor.trace.Value) !tensor.trace.Value {
1203     const product = try args[0].mul(args[1]);
1204     return try product.sum(.lane);
1205 }
1206 
1207 test "numerics: gradient of sum of products equals the other operand" {
1208     const allocator = std.testing.allocator;
1209     var source = try tensor.trace.define(allocator, "numerics_grad", &.{
1210         tensor.spec(.f32, .{ .lane = element_count }),
1211         tensor.spec(.f32, .{ .lane = element_count }),
1212     }, gradLoss);
1213     defer source.deinit();
1214 
1215     var differentiated = try tensor.grad(allocator, &source, .{ .wrt = &.{0} });
1216     defer differentiated.deinit();
1217 
1218     const module = try tensor.toSemanticModule(allocator, &differentiated);
1219 
1220     var a = [_]f32{ 1.0, -2.0, 3.0, 0.5, -0.25, 4.0, 7.0, -8.0 };
1221     var b = [_]f32{ 2.0, 0.5, -1.5, 8.0, 4.0, -0.125, 0.0, 1.0 };
1222     var out = @as([element_count]f32, @splat(0.0));
1223     var count = [_]i32{element_count};
1224 
1225     try runSingleKernel(allocator, module, &.{
1226         kernel.argumentBuffer(f32, out[0..]),
1227         kernel.argumentBuffer(f32, a[0..]),
1228         kernel.argumentBuffer(f32, b[0..]),
1229         kernel.argumentBuffer(i32, count[0..]),
1230     });
1231 
1232     try std.testing.expectEqualSlices(f32, b[0..], out[0..]);
1233 }
1234 
1235 test "numerics: one lowered kernel serves multiple effective extents through its count guard" {
1236     const allocator = std.testing.allocator;
1237     const module = try binaryF32Module(allocator, .add);
1238 
1239     const contract_module = try pipeline.prepareContractJobFromSemanticModule(allocator, module, .{});
1240     const tensor_module = try pipeline.prepareTensorJobFromContractJob(allocator, contract_module, .{});
1241     const dispatch_module = try pipeline.prepareDispatchJobFromTensorJob(allocator, tensor_module, .{});
1242     const memory_module = try pipeline.prepareMemoryJobFromDispatchJob(allocator, dispatch_module, .{});
1243     defer memory_module.deinit();
1244 
1245     var pass_ctx = memory_module.passContext();
1246     defer pass_ctx.deinit();
1247     const analysis = try kernelization.getKernelizationAnalysis(&pass_ctx, memory_module.choir_module);
1248     try std.testing.expectEqual(@as(usize, 1), analysis.kernelCount());
1249     const program = &analysis.kernels.items[0].program;
1250 
1251     var lhs = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 };
1252     var rhs = [_]f32{ 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0 };
1253 
1254     var full_out = @as([element_count]f32, @splat(0.0));
1255     var full_count = [_]i32{element_count};
1256     try runGeneratedProgram(allocator, program, &.{
1257         kernel.argumentBuffer(f32, full_out[0..]),
1258         kernel.argumentBuffer(f32, lhs[0..]),
1259         kernel.argumentBuffer(f32, rhs[0..]),
1260         kernel.argumentBuffer(i32, full_count[0..]),
1261     });
1262     try std.testing.expectEqualSlices(f32, &.{ 11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 77.0, 88.0 }, full_out[0..]);
1263 
1264     var short_out = @as([element_count]f32, @splat(0.0));
1265     var short_count = [_]i32{5};
1266     try runGeneratedProgram(allocator, program, &.{
1267         kernel.argumentBuffer(f32, short_out[0..]),
1268         kernel.argumentBuffer(f32, lhs[0..]),
1269         kernel.argumentBuffer(f32, rhs[0..]),
1270         kernel.argumentBuffer(i32, short_count[0..]),
1271     });
1272     try std.testing.expectEqualSlices(f32, &.{ 11.0, 22.0, 33.0, 44.0, 55.0, 0.0, 0.0, 0.0 }, short_out[0..]);
1273 }
1274 
1275 const IdentitySide = enum { left, right };
1276 
1277 fn matrixIdentityModule(
1278     allocator: std.mem.Allocator,
1279     side: IdentitySide,
1280     constant_identity: bool,
1281 ) !*SemanticModule {
1282     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1283     errdefer builder.deinit();
1284     const ty = try builder.tensor(.f32, &.{ 2, 2 });
1285     var fb = try builder.beginFunction(
1286         "matrix_identity",
1287         if (constant_identity) &.{ty} else &.{ ty, ty },
1288         &.{ty},
1289     );
1290     const payload = [_]f32{ 1, 0, 0, 1 };
1291     const identity = if (constant_identity)
1292         try fb.constant(ty, std.mem.sliceAsBytes(&payload))
1293     else
1294         fb.parameter(1);
1295     const lhs = if (side == .left) identity else fb.parameter(0);
1296     const rhs = if (side == .right) identity else fb.parameter(0);
1297     const product = try fb.dotGeneral(lhs, rhs, ty, &.{1}, &.{0}, &.{}, &.{});
1298     try fb.return_(&.{product});
1299     try fb.finish();
1300     return try builder.finish();
1301 }
1302 
1303 fn matrixIdentityReference(side: IdentitySide, input: [4]f32) [4]f32 {
1304     @setFloatMode(.strict);
1305     const identity = [_]f32{ 1, 0, 0, 1 };
1306     const lhs = if (side == .left) identity else input;
1307     const rhs = if (side == .right) identity else input;
1308     var output: [4]f32 = undefined;
1309     for (0..2) |row| {
1310         for (0..2) |column| {
1311             var sum: f32 = 0;
1312             for (0..2) |inner| sum += lhs[row * 2 + inner] * rhs[inner * 2 + column];
1313             output[row * 2 + column] = sum;
1314         }
1315     }
1316     return output;
1317 }
1318 
1319 fn expectF32Observation(expected: f32, actual: f32) !void {
1320     if (std.math.isNan(expected)) {
1321         try std.testing.expect(std.math.isNan(actual));
1322     } else {
1323         try std.testing.expectEqual(@as(u32, @bitCast(expected)), @as(u32, @bitCast(actual)));
1324     }
1325 }
1326 
1327 fn expectMatrixIdentityExecution(side: IdentitySide) !void {
1328     try @import("../fixture/root.zig").requireNativeCpuArtifacts();
1329     const allocator = std.testing.allocator;
1330     var state = gpu.cpu.State.init(allocator);
1331     defer state.deinit();
1332     const options = executable.FragmentCompilerOptions{ .artifact_format = .cpu_object };
1333     const folded = try compileAndLoadSemanticModule(
1334         allocator,
1335         state.handle(),
1336         try matrixIdentityModule(allocator, side, true),
1337         options,
1338     );
1339     defer folded.deinit();
1340     const runtime = try compileAndLoadSemanticModule(
1341         allocator,
1342         state.handle(),
1343         try matrixIdentityModule(allocator, side, false),
1344         options,
1345     );
1346     defer runtime.deinit();
1347     const cases = [_][4]f32{
1348         .{ 1.5, -2.25, 0.125, 42 },
1349         .{ std.math.inf(f32), 1, 2, 3 },
1350         .{ -std.math.inf(f32), 1, 2, 3 },
1351         .{ std.math.nan(f32), 1, 2, 3 },
1352         .{ -0.0, -0.0, -0.0, -0.0 },
1353         .{ 0.0, -0.0, -0.0, 0.0 },
1354     };
1355     const identity = [_]f32{ 1, 0, 0, 1 };
1356     for (cases) |input| {
1357         var folded_output: [4]f32 = undefined;
1358         var runtime_output: [4]f32 = undefined;
1359         const input_bytes = std.mem.sliceAsBytes(&input);
1360         try executable.invoke(folded, allocator, allocator, &.{input_bytes}, &.{
1361             std.mem.sliceAsBytes(&folded_output),
1362         });
1363         try executable.invoke(runtime, allocator, allocator, &.{
1364             input_bytes, std.mem.sliceAsBytes(&identity),
1365         }, &.{std.mem.sliceAsBytes(&runtime_output)});
1366         const expected = matrixIdentityReference(side, input);
1367         for (expected, runtime_output, folded_output) |reference, dynamic, constant| {
1368             try expectF32Observation(reference, dynamic);
1369             try expectF32Observation(reference, constant);
1370         }
1371     }
1372 }
1373 
1374 test "numerics: matrix identity on the left preserves floating observations" {
1375     try expectMatrixIdentityExecution(.left);
1376 }
1377 
1378 test "numerics: matrix identity on the right preserves floating observations" {
1379     try expectMatrixIdentityExecution(.right);
1380 }
1381 
1382 test "numerics: dot general matmul accumulates exactly on the oracle" {
1383     const allocator = std.testing.allocator;
1384     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1385     errdefer builder.deinit();
1386     const lhs_ty = try builder.tensor(.f32, &.{ 3, 4 });
1387     const rhs_ty = try builder.tensor(.f32, &.{ 4, 5 });
1388     const out_ty = try builder.tensor(.f32, &.{ 3, 5 });
1389     var fb = try builder.beginFunction("numerics_dot_general", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1390     const out_value = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{1}, &.{0}, &.{}, &.{});
1391     try fb.return_(&.{out_value});
1392     try fb.finish();
1393     const module = try builder.finish();
1394 
1395     var lhs: [12]f32 = undefined;
1396     for (&lhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index)) * 0.5 - 2.0;
1397     var rhs: [20]f32 = undefined;
1398     for (&rhs, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 7)) - 3.0;
1399 
1400     var expected: [15]f32 = undefined;
1401     for (0..3) |row| {
1402         for (0..5) |col| {
1403             var sum: f32 = 0.0;
1404             for (0..4) |inner| {
1405                 sum += lhs[row * 4 + inner] * rhs[inner * 5 + col];
1406             }
1407             expected[row * 5 + col] = sum;
1408         }
1409     }
1410 
1411     var out = @as([15]f32, @splat(0.0));
1412     try runSingleKernel(allocator, module, &.{
1413         kernel.argumentBuffer(f32, out[0..]),
1414         kernel.argumentBuffer(f32, lhs[0..]),
1415         kernel.argumentBuffer(f32, rhs[0..]),
1416     });
1417     try std.testing.expectEqualSlices(f32, expected[0..], out[0..]);
1418 }
1419 
1420 test "numerics: gather fuses into its consumer chain and matches the host" {
1421     const allocator = std.testing.allocator;
1422     const total = 512;
1423     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1424     errdefer builder.deinit();
1425     const vec_ty = try builder.tensor(.f32, &.{total});
1426     const idx_ty = try builder.tensor(.i32, &.{total});
1427     const scalar_ty = try builder.tensor(.f32, &.{});
1428     var fb = try builder.beginFunction("probe_gather_fusion", &.{vec_ty}, &.{vec_ty});
1429     const src = fb.parameter(0);
1430     const flat = try fb.iota(vec_ty, 0);
1431     const two = try fb.broadcast(try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, 2.0))), vec_ty, &.{total});
1432     const limit = try fb.broadcast(try fb.constant(scalar_ty, std.mem.asBytes(&@as(f32, total - 1))), vec_ty, &.{total});
1433     const doubled = try fb.min(try fb.mul(flat, two), limit);
1434     const idx = try fb.convert(doubled, idx_ty, .i32);
1435     const gathered = try fb.gather(src, idx, vec_ty, 0);
1436     const out_value = try fb.add(try fb.mul(gathered, two), flat);
1437     try fb.return_(&.{out_value});
1438     try fb.finish();
1439     const module = try builder.finish();
1440 
1441     var input: [total]f32 = undefined;
1442     for (input[0..], 0..) |*value, index| value.* = @sin(@as(f32, @floatFromInt(index)) * 0.05);
1443     var out = @as([total]f32, @splat(0.0));
1444     var count = [_]i32{total};
1445     try runSingleKernel(allocator, module, &.{
1446         kernel.argumentBuffer(f32, out[0..]),
1447         kernel.argumentBuffer(f32, input[0..]),
1448         kernel.argumentBuffer(i32, count[0..]),
1449     });
1450     for (0..total) |i| {
1451         const j = @min(i * 2, total - 1);
1452         const expected = input[j] * 2.0 + @as(f32, @floatFromInt(i));
1453         try std.testing.expectApproxEqAbs(expected, out[i], 1e-4);
1454     }
1455 }
1456 
1457 test "numerics: multistage tiled dot matches the host on CUDA" {
1458     const allocator = std.testing.allocator;
1459     var state = gpu.cuda.State.initDevice(allocator, 0) catch return error.SkipZigTest;
1460     defer state.deinit();
1461 
1462     const batch = 4;
1463     const m = 256;
1464     const n = 256;
1465     const k = 192;
1466     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1467     errdefer builder.deinit();
1468     const lhs_ty = try builder.tensor(.f32, &.{ batch, m, k });
1469     const rhs_ty = try builder.tensor(.f32, &.{ batch, k, n });
1470     const out_ty = try builder.tensor(.f32, &.{ batch, m, n });
1471     var fb = try builder.beginFunction("probe_multistage_dot", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1472     const out_value = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{2}, &.{1}, &.{0}, &.{0});
1473     try fb.return_(&.{out_value});
1474     try fb.finish();
1475     const module = try builder.finish();
1476 
1477     var fragment = try compileAndLoadSemanticModule(
1478         allocator,
1479         state.handle(),
1480         module,
1481         .{},
1482     );
1483     defer fragment.deinit();
1484 
1485     const lhs = try allocator.alloc(f32, batch * m * k);
1486     defer allocator.free(lhs);
1487     const rhs = try allocator.alloc(f32, batch * k * n);
1488     defer allocator.free(rhs);
1489     const out = try allocator.alloc(f32, batch * m * n);
1490     defer allocator.free(out);
1491     for (lhs, 0..) |*value, index| value.* = @sin(@as(f32, @floatFromInt(index)) * 0.013) * 0.5;
1492     for (rhs, 0..) |*value, index| value.* = @cos(@as(f32, @floatFromInt(index)) * 0.007) * 0.5;
1493     @memset(out, 0.0);
1494 
1495     var outputs = [_][]u8{std.mem.sliceAsBytes(out)};
1496     try executable.invoke(fragment, allocator, allocator, &.{
1497         std.mem.sliceAsBytes(lhs),
1498         std.mem.sliceAsBytes(rhs),
1499     }, &outputs);
1500 
1501     var sample: usize = 0;
1502     while (sample < 256) : (sample += 1) {
1503         const slab = sample % batch;
1504         const row = (sample * 37 + 11) % m;
1505         const col = (sample * 53 + 7) % n;
1506         var acc: f64 = 0.0;
1507         var inner: usize = 0;
1508         while (inner < k) : (inner += 1) {
1509             acc += @as(f64, lhs[(slab * m + row) * k + inner]) * @as(f64, rhs[(slab * k + inner) * n + col]);
1510         }
1511         const got: f64 = out[(slab * m + row) * n + col];
1512         try std.testing.expect(@abs(acc - got) < 1e-3);
1513     }
1514 }
1515 
1516 test "numerics: f16 dot general rounds through half precision" {
1517     const allocator = std.testing.allocator;
1518     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1519     errdefer builder.deinit();
1520     const lhs_ty = try builder.tensor(.f16, &.{ 3, 4 });
1521     const rhs_ty = try builder.tensor(.f16, &.{ 4, 5 });
1522     const out_ty = try builder.tensor(.f16, &.{ 3, 5 });
1523     var fb = try builder.beginFunction("numerics_f16_dot", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1524     const out_value = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{1}, &.{0}, &.{}, &.{});
1525     try fb.return_(&.{out_value});
1526     try fb.finish();
1527     const module = try builder.finish();
1528 
1529     var lhs = [_]f16{
1530         1.0,   -2.0, 3.0,  0.5,
1531         -0.25, 4.0,  7.0,  -8.0,
1532         0.75,  1.5,  -3.0, 2.0,
1533     };
1534     var rhs = [_]f16{
1535         2.0,    0.5,  -1.5, 8.0,  4.0,
1536         -0.125, 0.0,  1.0,  -2.0, 3.0,
1537         0.25,   -4.0, 2.0,  1.0,  -1.0,
1538         3.0,    2.0,  0.5,  -0.5, 1.5,
1539     };
1540     var expected = @as([15]f16, @splat(0.0));
1541     for (0..3) |row| {
1542         for (0..5) |col| {
1543             var sum: f16 = 0.0;
1544             for (0..4) |inner| {
1545                 const product: f16 = @floatCast(@as(f32, lhs[row * 4 + inner]) * @as(f32, rhs[inner * 5 + col]));
1546                 sum = @floatCast(@as(f32, sum) + @as(f32, product));
1547             }
1548             expected[row * 5 + col] = sum;
1549         }
1550     }
1551 
1552     var out = @as([15]f16, @splat(0.0));
1553     try runSingleKernel(allocator, module, &.{
1554         kernel.argumentBuffer(f16, out[0..]),
1555         kernel.argumentBuffer(f16, lhs[0..]),
1556         kernel.argumentBuffer(f16, rhs[0..]),
1557     });
1558     for (expected, out) |want, actual| {
1559         try std.testing.expectApproxEqAbs(@as(f32, want), @as(f32, actual), 0.001);
1560     }
1561 }
1562 
1563 test "numerics: i32 arithmetic wraps on overflow like the host" {
1564     const allocator = std.testing.allocator;
1565     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1566     errdefer builder.deinit();
1567     const vec_ty = try builder.tensor(.i32, &.{element_count});
1568     var fb = try builder.beginFunction("numerics_i32_overflow", &.{ vec_ty, vec_ty }, &.{vec_ty});
1569     const out_value = try fb.add(fb.parameter(0), fb.parameter(1));
1570     try fb.return_(&.{out_value});
1571     try fb.finish();
1572     const module = try builder.finish();
1573 
1574     const max = std.math.maxInt(i32);
1575     const min = std.math.minInt(i32);
1576     var lhs = [_]i32{ max, min, max, min, 1, -1, max - 1, min + 1 };
1577     var rhs = [_]i32{ 1, -1, max, min, max, min, 2, -2 };
1578     var out = @as([element_count]i32, @splat(0));
1579     var count = [_]i32{element_count};
1580 
1581     try runSingleKernel(allocator, module, &.{
1582         kernel.argumentBuffer(i32, out[0..]),
1583         kernel.argumentBuffer(i32, lhs[0..]),
1584         kernel.argumentBuffer(i32, rhs[0..]),
1585         kernel.argumentBuffer(i32, count[0..]),
1586     });
1587 
1588     for (lhs, rhs, out) |a, b, actual| {
1589         try std.testing.expectEqual(a +% b, actual);
1590     }
1591 }
1592 
1593 test "numerics: f16 add stays within half precision tolerance" {
1594     const allocator = std.testing.allocator;
1595     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1596     errdefer builder.deinit();
1597     const vec_ty = try builder.tensor(.f16, &.{element_count});
1598     var fb = try builder.beginFunction("numerics_f16_add", &.{ vec_ty, vec_ty }, &.{vec_ty});
1599     const out_value = try fb.add(fb.parameter(0), fb.parameter(1));
1600     try fb.return_(&.{out_value});
1601     try fb.finish();
1602     const module = try builder.finish();
1603 
1604     var lhs = [_]f16{ 1.0, -2.5, 3.25, 0.0, 100.0, -0.125, 7.5, -8.0 };
1605     var rhs = [_]f16{ 0.5, 2.5, -3.25, -0.0, 1.0, 0.125, 0.25, 8.0 };
1606     var out = @as([element_count]f16, @splat(0.0));
1607     var count = [_]i32{element_count};
1608 
1609     try runSingleKernel(allocator, module, &.{
1610         kernel.argumentBuffer(f16, out[0..]),
1611         kernel.argumentBuffer(f16, lhs[0..]),
1612         kernel.argumentBuffer(f16, rhs[0..]),
1613         kernel.argumentBuffer(i32, count[0..]),
1614     });
1615 
1616     for (lhs, rhs, out) |a, b, actual| {
1617         const expected: f32 = @as(f32, a) + @as(f32, b);
1618         try std.testing.expectApproxEqAbs(expected, @as(f32, actual), 0.001);
1619     }
1620 }
1621 
1622 test "numerics: f16 tanh stays within half precision tolerance" {
1623     const allocator = std.testing.allocator;
1624     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1625     errdefer builder.deinit();
1626     const vec_ty = try builder.tensor(.f16, &.{element_count});
1627     var fb = try builder.beginFunction("numerics_f16_tanh", &.{vec_ty}, &.{vec_ty});
1628     const out_value = try fb.tanh(fb.parameter(0));
1629     try fb.return_(&.{out_value});
1630     try fb.finish();
1631     const module = try builder.finish();
1632 
1633     var input = [_]f16{ -4.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };
1634     var out = @as([element_count]f16, @splat(0.0));
1635     var count = [_]i32{element_count};
1636 
1637     try runSingleKernel(allocator, module, &.{
1638         kernel.argumentBuffer(f16, out[0..]),
1639         kernel.argumentBuffer(f16, input[0..]),
1640         kernel.argumentBuffer(i32, count[0..]),
1641     });
1642 
1643     for (input, out) |x, actual| {
1644         const expected: f32 = std.math.tanh(@as(f32, x));
1645         try std.testing.expectApproxEqAbs(expected, @as(f32, actual), 0.001);
1646     }
1647 }
1648 
1649 const ReturnedView = enum { broadcast, broadcast_in_dim, reshape_broadcast, computed_broadcast };
1650 
1651 fn returnedViewModule(allocator: std.mem.Allocator, view: ReturnedView) !*SemanticModule {
1652     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1653     errdefer builder.deinit();
1654     const input_dims: []const i64 = switch (view) {
1655         .broadcast => &.{},
1656         .broadcast_in_dim => &.{2},
1657         .reshape_broadcast, .computed_broadcast => &.{4},
1658     };
1659     const input_type = try builder.tensor(.f32, input_dims);
1660     const broadcast_type = try builder.tensor(.f32, &.{ 2, 4 });
1661     const output_type = if (view == .reshape_broadcast or view == .computed_broadcast)
1662         try builder.tensor(.f32, &.{8})
1663     else
1664         broadcast_type;
1665     var function = try builder.beginFunction(@tagName(view), &.{input_type}, &.{output_type});
1666     const output = switch (view) {
1667         .broadcast => try function.broadcast(function.parameter(0), output_type, &.{ 2, 4 }),
1668         .broadcast_in_dim => try function.broadcastInDim(
1669             function.parameter(0),
1670             output_type,
1671             &.{ 2, 4 },
1672             &.{0},
1673         ),
1674         .reshape_broadcast => try function.reshape(
1675             try function.broadcast(function.parameter(0), broadcast_type, &.{2}),
1676             output_type,
1677             &.{8},
1678         ),
1679         .computed_broadcast => try function.reshape(
1680             try function.broadcast(
1681                 try function.mul(function.parameter(0), function.parameter(0)),
1682                 broadcast_type,
1683                 &.{2},
1684             ),
1685             output_type,
1686             &.{8},
1687         ),
1688     };
1689     try function.return_(&.{output});
1690     try function.finish();
1691     return try builder.finish();
1692 }
1693 
1694 fn expectReturnedViewReadback(view: ReturnedView) !void {
1695     try @import("../fixture/root.zig").requireNativeCpuArtifacts();
1696     const allocator = std.testing.allocator;
1697     var state = gpu.cpu.State.init(allocator);
1698     defer state.deinit();
1699     const fragment = try compileAndLoadSemanticModule(
1700         allocator,
1701         state.handle(),
1702         try returnedViewModule(allocator, view),
1703         .{ .artifact_format = .cpu_object },
1704     );
1705     defer fragment.deinit();
1706     const inputs = [_]f32{ 7.5, -2.25, 0.125, 42 };
1707     const count: usize = switch (view) {
1708         .broadcast => 1,
1709         .broadcast_in_dim => 2,
1710         .reshape_broadcast, .computed_broadcast => 4,
1711     };
1712     var invocation = try executable.prepareInvocation(fragment, allocator, &.{
1713         std.mem.sliceAsBytes(inputs[0..count]),
1714     });
1715     defer invocation.deinit();
1716     try invocation.launch(allocator);
1717     var output: [8]f32 = undefined;
1718     try invocation.readOutput(0, std.mem.sliceAsBytes(&output));
1719     for (output, 0..) |actual, index| {
1720         const source = switch (view) {
1721             .broadcast => 0,
1722             .broadcast_in_dim => index / 4,
1723             .reshape_broadcast, .computed_broadcast => index % 4,
1724         };
1725         const expected = if (view == .computed_broadcast)
1726             inputs[source] * inputs[source]
1727         else
1728             inputs[source];
1729         try expectF32Observation(expected, actual);
1730     }
1731 }
1732 
1733 test "numerics: returned standalone broadcast has an executed readback" {
1734     try expectReturnedViewReadback(.broadcast);
1735 }
1736 
1737 test "numerics: returned standalone broadcast_in_dim has an executed readback" {
1738     try expectReturnedViewReadback(.broadcast_in_dim);
1739 }
1740 
1741 test "numerics: returned reshape of broadcast has an executed readback" {
1742     try expectReturnedViewReadback(.reshape_broadcast);
1743 }
1744 
1745 test "numerics: returned reshape of computed broadcast respects writer dependencies" {
1746     try expectReturnedViewReadback(.computed_broadcast);
1747 }
1748 
1749 test "numerics: returned broadcast beyond constant folding limit has an executed readback" {
1750     try @import("../fixture/root.zig").requireNativeCpuArtifacts();
1751     const allocator = std.testing.allocator;
1752     const count = 16385;
1753     const expected: f32 = -2.25;
1754     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1755     defer builder.deinit();
1756     const scalar = try builder.tensor(.f32, &.{});
1757     const vector = try builder.tensor(.f32, &.{count});
1758     var function = try builder.beginFunction("large_constant_broadcast", &.{}, &.{vector});
1759     const constant = try function.constant(scalar, std.mem.asBytes(&expected));
1760     const broadcast = try function.broadcast(constant, vector, &.{count});
1761     try function.return_(&.{broadcast});
1762     try function.finish();
1763     var state = gpu.cpu.State.init(allocator);
1764     defer state.deinit();
1765     const fragment = try compileAndLoadSemanticModule(
1766         allocator,
1767         state.handle(),
1768         try builder.finish(),
1769         .{ .artifact_format = .cpu_object },
1770     );
1771     defer fragment.deinit();
1772     try std.testing.expectEqual(@as(usize, 1), fragment.kernelCount());
1773     var invocation = try executable.prepareInvocation(fragment, allocator, &.{});
1774     defer invocation.deinit();
1775     try invocation.launch(allocator);
1776     const output = try allocator.alloc(f32, count);
1777     defer allocator.free(output);
1778     try invocation.readOutput(0, std.mem.sliceAsBytes(output));
1779     for (output) |actual| try expectF32Observation(expected, actual);
1780 }
1781 
1782 test "numerics: returned input and constant backing have executed readbacks without kernels" {
1783     try @import("../fixture/root.zig").requireNativeCpuArtifacts();
1784     const allocator = std.testing.allocator;
1785     const input = [_]f32{ 7.5, -2.25 };
1786     const payload = [_]f32{ 0.125, 42 };
1787     var builder = try SemanticBuilder.init(allocator, SemanticBuilder.ContextLimits.standard);
1788     defer builder.deinit();
1789     const vector = try builder.tensor(.f32, &.{2});
1790     var function = try builder.beginFunction("backed_outputs", &.{vector}, &.{ vector, vector });
1791     const constant = try function.constant(vector, std.mem.sliceAsBytes(&payload));
1792     try function.return_(&.{ function.parameter(0), constant });
1793     try function.finish();
1794     var state = gpu.cpu.State.init(allocator);
1795     defer state.deinit();
1796     const fragment = try compileAndLoadSemanticModule(
1797         allocator,
1798         state.handle(),
1799         try builder.finish(),
1800         .{ .artifact_format = .cpu_object },
1801     );
1802     defer fragment.deinit();
1803     try std.testing.expectEqual(@as(usize, 0), fragment.kernelCount());
1804     var invocation = try executable.prepareInvocation(fragment, allocator, &.{
1805         std.mem.sliceAsBytes(&input),
1806     });
1807     defer invocation.deinit();
1808     try invocation.launch(allocator);
1809     for ([_][2]f32{ input, payload }, 0..) |expected, index| {
1810         var output: [2]f32 = undefined;
1811         try invocation.readOutput(index, std.mem.sliceAsBytes(&output));
1812         for (expected, output) |wanted, actual| try expectF32Observation(wanted, actual);
1813     }
1814 }