lib/accy/src/choir/semantics.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const testing = std.testing;
4 const axis_roles = @import("../axis/root.zig").roles;
5 const activation = @import("activation.zig");
6 const einsum = @import("einsum/root.zig");
7
8 pub const Allocator = std.mem.Allocator;
9 pub const DType = choir_abi.DType;
10
11 pub const Type = struct {
12 dtype: DType,
13 dims: []const i64,
14
15 pub fn eql(a: Type, b: Type) bool {
16 return a.dtype == b.dtype and std.mem.eql(i64, a.dims, b.dims);
17 }
18 };
19
20 pub const OpKind = enum(u8) {
21 constant,
22 parameter,
23 iota,
24 add,
25 sub,
26 mul,
27 div,
28 max,
29 min,
30 pow,
31 compare,
32 neg,
33 exp,
34 log,
35 tanh,
36 sqrt,
37 activation,
38 abs,
39 sin,
40 cos,
41 tan,
42 floor,
43 round,
44 trunc,
45 atan2,
46 convert,
47 reduce,
48 dot_general,
49 einsum,
50 broadcast,
51 broadcast_in_dim,
52 reshape,
53 transpose,
54 slice,
55 gather,
56 scatter,
57 scatter_add,
58 sparse_cross_entropy,
59 pad,
60 concatenate,
61 select,
62 cumsum,
63 scratch,
64 kernel_call,
65 iterate,
66 iterate_yield,
67 @"return",
68 };
69
70 pub const CompareDirection = enum { eq, ne, lt, le, gt, ge };
71
72 pub const ActivationKind = activation.Kind;
73
74 pub const ReducerKind = enum { sum, max, min };
75
76 pub const KernelOperandEffect = enum(u8) {
77 none,
78 read,
79 write,
80 read_write,
81 unknown,
82
83 pub fn fromByte(value: u8) ?KernelOperandEffect {
84 return switch (value) {
85 @backingInt(KernelOperandEffect.none) => .none,
86 @backingInt(KernelOperandEffect.read) => .read,
87 @backingInt(KernelOperandEffect.write) => .write,
88 @backingInt(KernelOperandEffect.read_write) => .read_write,
89 @backingInt(KernelOperandEffect.unknown) => .unknown,
90 else => null,
91 };
92 }
93
94 pub fn writes(self: KernelOperandEffect) bool {
95 return switch (self) {
96 .write, .read_write, .unknown => true,
97 .none, .read => false,
98 };
99 }
100 };
101
102 pub const KernelCallContract = struct {
103 target: []const u8,
104 version: u32,
105 has_side_effects: bool,
106 operand_effects: []const KernelOperandEffect,
107 result_aliases: []const ?usize,
108 results: []const Type,
109 };
110
111 pub const Attribute = union(enum) {
112 i64: i64,
113 i64_list: []const i64,
114 dtype: DType,
115 compare_direction: CompareDirection,
116 activation_kind: ActivationKind,
117 reducer_kind: ReducerKind,
118 kernel_call: KernelCallContract,
119 bytes: []const u8,
120 einsum: []const u8,
121 };
122
123 pub const InferError = error{
124 ArityMismatch,
125 DTypeMismatch,
126 ShapeMismatch,
127 RankMismatch,
128 AttributeMissing,
129 AttributeKindMismatch,
130 InvalidDimension,
131 ElementCountMismatch,
132 InvalidPermutation,
133 DimMismatch,
134 ContractionMismatch,
135 InvalidEinsum,
136 InvalidKernelContract,
137 } || Allocator.Error;
138
139 pub const InferFn = *const fn (
140 arena: Allocator,
141 inputs: []const Type,
142 attrs: []const Attribute,
143 ) InferError![]const Type;
144
145 pub const SupportedDTypes = enum {
146 any,
147 arithmetic,
148 float_only,
149
150 pub fn allows(self: SupportedDTypes, dt: DType) bool {
151 return switch (self) {
152 .any => true,
153 .arithmetic => dt.isNumeric(),
154 .float_only => dt.isFloat(),
155 };
156 }
157 };
158
159 pub const OpInfo = struct {
160 kind: OpKind,
161 name: []const u8,
162 supported_dtypes: SupportedDTypes,
163 infer: InferFn,
164 };
165
166 pub fn info(kind: OpKind) OpInfo {
167 return switch (kind) {
168 .constant => .{ .kind = .constant, .name = "constant", .supported_dtypes = .any, .infer = inferConstant },
169 .parameter => .{ .kind = .parameter, .name = "parameter", .supported_dtypes = .any, .infer = inferParameter },
170 .iota => .{ .kind = .iota, .name = "iota", .supported_dtypes = .any, .infer = inferIota },
171 .add => .{ .kind = .add, .name = "add", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
172 .sub => .{ .kind = .sub, .name = "sub", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
173 .mul => .{ .kind = .mul, .name = "mul", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
174 .div => .{ .kind = .div, .name = "div", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
175 .max => .{ .kind = .max, .name = "max", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
176 .min => .{ .kind = .min, .name = "min", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },
177 .pow => .{ .kind = .pow, .name = "pow", .supported_dtypes = .float_only, .infer = inferElementwiseBinary },
178 .compare => .{ .kind = .compare, .name = "compare", .supported_dtypes = .any, .infer = inferCompare },
179 .neg => .{ .kind = .neg, .name = "neg", .supported_dtypes = .arithmetic, .infer = inferElementwiseUnary },
180 .exp => .{ .kind = .exp, .name = "exp", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
181 .log => .{ .kind = .log, .name = "log", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
182 .tanh => .{ .kind = .tanh, .name = "tanh", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
183 .sqrt => .{ .kind = .sqrt, .name = "sqrt", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
184 .activation => .{ .kind = .activation, .name = "activation", .supported_dtypes = .float_only, .infer = inferActivation },
185 .abs => .{ .kind = .abs, .name = "abs", .supported_dtypes = .arithmetic, .infer = inferElementwiseUnary },
186 .sin => .{ .kind = .sin, .name = "sin", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
187 .cos => .{ .kind = .cos, .name = "cos", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
188 .tan => .{ .kind = .tan, .name = "tan", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
189 .floor => .{ .kind = .floor, .name = "floor", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
190 .round => .{ .kind = .round, .name = "round", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
191 .trunc => .{ .kind = .trunc, .name = "trunc", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },
192 .atan2 => .{ .kind = .atan2, .name = "atan2", .supported_dtypes = .float_only, .infer = inferElementwiseBinary },
193 .convert => .{ .kind = .convert, .name = "convert", .supported_dtypes = .any, .infer = inferConvert },
194 .reduce => .{ .kind = .reduce, .name = "reduce", .supported_dtypes = .any, .infer = inferReduce },
195 .dot_general => .{ .kind = .dot_general, .name = "dot_general", .supported_dtypes = .arithmetic, .infer = inferDotGeneral },
196 .einsum => .{ .kind = .einsum, .name = "einsum", .supported_dtypes = .arithmetic, .infer = inferEinsum },
197 .broadcast => .{ .kind = .broadcast, .name = "broadcast", .supported_dtypes = .any, .infer = inferBroadcast },
198 .broadcast_in_dim => .{ .kind = .broadcast_in_dim, .name = "broadcast_in_dim", .supported_dtypes = .any, .infer = inferBroadcastInDim },
199 .reshape => .{ .kind = .reshape, .name = "reshape", .supported_dtypes = .any, .infer = inferReshape },
200 .transpose => .{ .kind = .transpose, .name = "transpose", .supported_dtypes = .any, .infer = inferTranspose },
201 .slice => .{ .kind = .slice, .name = "slice", .supported_dtypes = .any, .infer = inferSlice },
202 .gather => .{ .kind = .gather, .name = "gather", .supported_dtypes = .any, .infer = inferGather },
203 .scatter => .{ .kind = .scatter, .name = "scatter", .supported_dtypes = .any, .infer = inferScatter },
204 .scatter_add => .{ .kind = .scatter_add, .name = "scatter_add", .supported_dtypes = .any, .infer = inferScatter },
205 .sparse_cross_entropy => .{ .kind = .sparse_cross_entropy, .name = "sparse_cross_entropy", .supported_dtypes = .any, .infer = inferSparseCrossEntropy },
206 .pad => .{ .kind = .pad, .name = "pad", .supported_dtypes = .any, .infer = inferPad },
207 .concatenate => .{ .kind = .concatenate, .name = "concatenate", .supported_dtypes = .any, .infer = inferConcatenate },
208 .select => .{ .kind = .select, .name = "select", .supported_dtypes = .any, .infer = inferSelect },
209 .cumsum => .{ .kind = .cumsum, .name = "cumsum", .supported_dtypes = .arithmetic, .infer = inferCumsum },
210 .scratch => .{ .kind = .scratch, .name = "scratch", .supported_dtypes = .any, .infer = inferScratch },
211 .kernel_call => .{ .kind = .kernel_call, .name = "kernel_call", .supported_dtypes = .any, .infer = inferKernelCall },
212 .iterate => .{ .kind = .iterate, .name = "iterate", .supported_dtypes = .any, .infer = inferIterate },
213 .iterate_yield => .{ .kind = .iterate_yield, .name = "iterate_yield", .supported_dtypes = .any, .infer = inferIterateYield },
214 .@"return" => .{ .kind = .@"return", .name = "return", .supported_dtypes = .any, .infer = inferReturn },
215 };
216 }
217
218 pub fn inferShape(
219 kind: OpKind,
220 arena: Allocator,
221 inputs: []const Type,
222 attrs: []const Attribute,
223 ) InferError![]const Type {
224 return try info(kind).infer(arena, inputs, attrs);
225 }
226
227 fn inferAttrI64(attrs: []const Attribute, idx: usize) InferError!i64 {
228 if (idx >= attrs.len) return error.AttributeMissing;
229 return switch (attrs[idx]) {
230 .i64 => |v| v,
231 else => error.AttributeKindMismatch,
232 };
233 }
234
235 fn inferAttrI64List(attrs: []const Attribute, idx: usize) InferError![]const i64 {
236 if (idx >= attrs.len) return error.AttributeMissing;
237 return switch (attrs[idx]) {
238 .i64_list => |v| v,
239 else => error.AttributeKindMismatch,
240 };
241 }
242
243 fn inferAttrDType(attrs: []const Attribute, idx: usize) InferError!DType {
244 if (idx >= attrs.len) return error.AttributeMissing;
245 return switch (attrs[idx]) {
246 .dtype => |v| v,
247 else => error.AttributeKindMismatch,
248 };
249 }
250
251 fn inferAttrCompareDirection(attrs: []const Attribute, idx: usize) InferError!CompareDirection {
252 if (idx >= attrs.len) return error.AttributeMissing;
253 return switch (attrs[idx]) {
254 .compare_direction => |v| v,
255 else => error.AttributeKindMismatch,
256 };
257 }
258
259 fn inferAttrActivationKind(attrs: []const Attribute, idx: usize) InferError!ActivationKind {
260 if (idx >= attrs.len) return error.AttributeMissing;
261 return switch (attrs[idx]) {
262 .activation_kind => |v| v,
263 else => error.AttributeKindMismatch,
264 };
265 }
266
267 fn inferAttrReducerKind(attrs: []const Attribute, idx: usize) InferError!ReducerKind {
268 if (idx >= attrs.len) return error.AttributeMissing;
269 return switch (attrs[idx]) {
270 .reducer_kind => |v| v,
271 else => error.AttributeKindMismatch,
272 };
273 }
274
275 fn inferAttrBytes(attrs: []const Attribute, idx: usize) InferError![]const u8 {
276 if (idx >= attrs.len) return error.AttributeMissing;
277 return switch (attrs[idx]) {
278 .bytes => |v| v,
279 else => error.AttributeKindMismatch,
280 };
281 }
282
283 fn inferAttrEinsum(attrs: []const Attribute, idx: usize) InferError![]const u8 {
284 if (idx >= attrs.len) return error.AttributeMissing;
285 return switch (attrs[idx]) {
286 .einsum => |v| v,
287 else => error.AttributeKindMismatch,
288 };
289 }
290
291 fn inferAttrKernelCall(attrs: []const Attribute, idx: usize) InferError!KernelCallContract {
292 if (idx >= attrs.len) return error.AttributeMissing;
293 return switch (attrs[idx]) {
294 .kernel_call => |v| v,
295 else => error.AttributeKindMismatch,
296 };
297 }
298
299 fn singletonType(arena: Allocator, t: Type) InferError![]const Type {
300 const out = try arena.alloc(Type, 1);
301 out[0] = t;
302 return out;
303 }
304
305 fn duplicateDims(arena: Allocator, dims: []const i64) InferError![]const i64 {
306 return arena.dupe(i64, dims);
307 }
308
309 fn duplicateTypes(arena: Allocator, types: []const Type) InferError![]const Type {
310 const out = try arena.alloc(Type, types.len);
311 for (types, 0..) |typ, i| {
312 for (typ.dims) |d| if (d < 0) return error.InvalidDimension;
313 out[i] = .{
314 .dtype = typ.dtype,
315 .dims = try duplicateDims(arena, typ.dims),
316 };
317 }
318 return out;
319 }
320
321 pub fn inferConstant(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
322 if (inputs.len != 0) return error.ArityMismatch;
323 _ = try inferAttrBytes(attrs, 0);
324 const dt = try inferAttrDType(attrs, 1);
325 const dims = try inferAttrI64List(attrs, 2);
326 return singletonType(arena, .{
327 .dtype = dt,
328 .dims = try duplicateDims(arena, dims),
329 });
330 }
331
332 pub fn inferParameter(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
333 if (inputs.len != 0) return error.ArityMismatch;
334 const idx = try inferAttrI64(attrs, 0);
335 if (idx < 0) return error.InvalidDimension;
336 const dt = try inferAttrDType(attrs, 1);
337 const dims = try inferAttrI64List(attrs, 2);
338 return singletonType(arena, .{
339 .dtype = dt,
340 .dims = try duplicateDims(arena, dims),
341 });
342 }
343
344 pub fn inferIota(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
345 if (inputs.len != 0) return error.ArityMismatch;
346 const dim = try inferAttrI64(attrs, 0);
347 const dt = try inferAttrDType(attrs, 1);
348 const dims = try inferAttrI64List(attrs, 2);
349 if (dim < 0 or dim >= @as(i64, @intCast(dims.len))) return error.InvalidDimension;
350 return singletonType(arena, .{
351 .dtype = dt,
352 .dims = try duplicateDims(arena, dims),
353 });
354 }
355
356 pub fn inferElementwiseBinary(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
357 _ = attrs;
358 if (inputs.len != 2) return error.ArityMismatch;
359 if (inputs[0].dtype != inputs[1].dtype) return error.DTypeMismatch;
360 if (!std.mem.eql(i64, inputs[0].dims, inputs[1].dims)) return error.ShapeMismatch;
361 return singletonType(arena, .{
362 .dtype = inputs[0].dtype,
363 .dims = try duplicateDims(arena, inputs[0].dims),
364 });
365 }
366
367 pub fn inferCompare(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
368 _ = try inferAttrCompareDirection(attrs, 0);
369 if (inputs.len != 2) return error.ArityMismatch;
370 if (inputs[0].dtype != inputs[1].dtype) return error.DTypeMismatch;
371 if (!std.mem.eql(i64, inputs[0].dims, inputs[1].dims)) return error.ShapeMismatch;
372 return singletonType(arena, .{
373 .dtype = .i1,
374 .dims = try duplicateDims(arena, inputs[0].dims),
375 });
376 }
377
378 pub fn inferElementwiseUnary(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
379 _ = attrs;
380 if (inputs.len != 1) return error.ArityMismatch;
381 return singletonType(arena, .{
382 .dtype = inputs[0].dtype,
383 .dims = try duplicateDims(arena, inputs[0].dims),
384 });
385 }
386
387 pub fn inferActivation(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
388 _ = try inferAttrActivationKind(attrs, 0);
389 return inferElementwiseUnary(arena, inputs, attrs);
390 }
391
392 pub fn inferConvert(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
393 if (inputs.len != 1) return error.ArityMismatch;
394 const target = try inferAttrDType(attrs, 0);
395 return singletonType(arena, .{
396 .dtype = target,
397 .dims = try duplicateDims(arena, inputs[0].dims),
398 });
399 }
400
401 pub fn inferReduce(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
402 if (inputs.len != 2) return error.ArityMismatch;
403 _ = try inferAttrReducerKind(attrs, 0);
404 const dims = try inferAttrI64List(attrs, 1);
405 const operand = inputs[0];
406 const init = inputs[1];
407 if (operand.dtype != init.dtype) return error.DTypeMismatch;
408 if (init.dims.len != 0) return error.RankMismatch;
409 const rank: i64 = @intCast(operand.dims.len);
410 for (dims, 0..) |d, i| {
411 if (d < 0 or d >= rank) return error.InvalidDimension;
412 for (dims[0..i]) |e| if (e == d) return error.InvalidDimension;
413 }
414 const out_rank = operand.dims.len - dims.len;
415 const out_dims = try arena.alloc(i64, out_rank);
416 var j: usize = 0;
417 outer: for (operand.dims, 0..) |size, idx_u| {
418 const idx: i64 = @intCast(idx_u);
419 for (dims) |d| if (d == idx) continue :outer;
420 out_dims[j] = size;
421 j += 1;
422 }
423 std.debug.assert(j == out_rank);
424 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
425 }
426
427 pub fn inferDotGeneral(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
428 if (inputs.len != 2) return error.ArityMismatch;
429 const lhs = inputs[0];
430 const rhs = inputs[1];
431 if (lhs.dtype != rhs.dtype) return error.DTypeMismatch;
432 const lhs_batch = try inferAttrI64List(attrs, 0);
433 const rhs_batch = try inferAttrI64List(attrs, 1);
434 const lhs_contract = try inferAttrI64List(attrs, 2);
435 const rhs_contract = try inferAttrI64List(attrs, 3);
436 const result_dtype = try inferDotGeneralResultDType(lhs.dtype, attrs);
437 if (lhs_batch.len != rhs_batch.len) return error.ContractionMismatch;
438 if (lhs_contract.len != rhs_contract.len) return error.ContractionMismatch;
439
440 if (axis_roles.check(lhs.dims.len, &.{ lhs_batch, lhs_contract }) != null) {
441 return error.InvalidDimension;
442 }
443 if (axis_roles.check(rhs.dims.len, &.{ rhs_batch, rhs_contract }) != null) {
444 return error.InvalidDimension;
445 }
446
447 for (lhs_batch, rhs_batch) |ld, rd| {
448 if (lhs.dims[@intCast(ld)] != rhs.dims[@intCast(rd)]) return error.ContractionMismatch;
449 }
450 for (lhs_contract, rhs_contract) |ld, rd| {
451 if (lhs.dims[@intCast(ld)] != rhs.dims[@intCast(rd)]) return error.ContractionMismatch;
452 }
453
454 const out_rank = lhs_batch.len +
455 (lhs.dims.len - lhs_batch.len - lhs_contract.len) +
456 (rhs.dims.len - rhs_batch.len - rhs_contract.len);
457 const out_dims = try arena.alloc(i64, out_rank);
458 var w: usize = 0;
459 for (lhs_batch) |d| {
460 out_dims[w] = lhs.dims[@intCast(d)];
461 w += 1;
462 }
463 for (lhs.dims, 0..) |size, idx_u| {
464 const idx: i64 = @intCast(idx_u);
465 if (containsDim(lhs_batch, idx)) continue;
466 if (containsDim(lhs_contract, idx)) continue;
467 out_dims[w] = size;
468 w += 1;
469 }
470 for (rhs.dims, 0..) |size, idx_u| {
471 const idx: i64 = @intCast(idx_u);
472 if (containsDim(rhs_batch, idx)) continue;
473 if (containsDim(rhs_contract, idx)) continue;
474 out_dims[w] = size;
475 w += 1;
476 }
477 std.debug.assert(w == out_rank);
478 return singletonType(arena, .{ .dtype = result_dtype, .dims = out_dims });
479 }
480
481 pub fn inferEinsum(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
482 if (inputs.len == 0) return error.ArityMismatch;
483 const equation_text = try inferAttrEinsum(attrs, 0);
484 const dtype = inputs[0].dtype;
485 const shapes = try arena.alloc([]const u64, inputs.len);
486 for (inputs, 0..) |input, input_index| {
487 if (input.dtype != dtype) return error.DTypeMismatch;
488 const dims = try arena.alloc(u64, input.dims.len);
489 for (input.dims, 0..) |dim, dim_index| {
490 if (dim < 0) return error.InvalidDimension;
491 dims[dim_index] = @intCast(dim);
492 }
493 shapes[input_index] = dims;
494 }
495
496 var parsed = einsum.parse(arena, equation_text, shapes) catch return error.InvalidEinsum;
497 defer parsed.deinit();
498
499 const out_dims = try arena.alloc(i64, parsed.output.len);
500 for (parsed.output, 0..) |index, i| {
501 const dim = parsed.dimension(index);
502 if (dim > @as(u64, @intCast(std.math.maxInt(i64)))) return error.InvalidDimension;
503 out_dims[i] = @intCast(dim);
504 }
505 return singletonType(arena, .{ .dtype = dtype, .dims = out_dims });
506 }
507
508 fn inferDotGeneralResultDType(input_dtype: DType, attrs: []const Attribute) InferError!DType {
509 if (attrs.len == 4) return input_dtype;
510 if (attrs.len != 5) return error.AttributeKindMismatch;
511 const result_dtype = try inferAttrDType(attrs, 4);
512 if (!dotGeneralResultDTypeAllowed(input_dtype, result_dtype)) return error.DTypeMismatch;
513 return result_dtype;
514 }
515
516 fn dotGeneralResultDTypeAllowed(input_dtype: DType, result_dtype: DType) bool {
517 return input_dtype == result_dtype or (input_dtype == .f16 and result_dtype == .f32);
518 }
519
520 fn containsDim(list: []const i64, d: i64) bool {
521 for (list) |e| if (e == d) return true;
522 return false;
523 }
524
525 pub fn inferBroadcast(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
526 if (inputs.len != 1) return error.ArityMismatch;
527 const sizes = try inferAttrI64List(attrs, 0);
528 for (sizes) |s| if (s < 0) return error.InvalidDimension;
529 const out_dims = try arena.alloc(i64, sizes.len + inputs[0].dims.len);
530 @memcpy(out_dims[0..sizes.len], sizes);
531 @memcpy(out_dims[sizes.len..], inputs[0].dims);
532 return singletonType(arena, .{ .dtype = inputs[0].dtype, .dims = out_dims });
533 }
534
535 pub fn inferBroadcastInDim(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
536 if (inputs.len != 1) return error.ArityMismatch;
537 const broadcast_dims = try inferAttrI64List(attrs, 0);
538 const result_shape = try inferAttrI64List(attrs, 1);
539 const operand = inputs[0];
540
541 if (broadcast_dims.len != operand.dims.len) return error.ArityMismatch;
542 const result_rank: i64 = @intCast(result_shape.len);
543 for (broadcast_dims, 0..) |d, i| {
544 if (d < 0 or d >= result_rank) return error.InvalidDimension;
545 for (broadcast_dims[0..i]) |e| if (e == d) return error.InvalidDimension;
546 }
547 for (broadcast_dims, operand.dims) |target, operand_dim| {
548 const out_dim = result_shape[@intCast(target)];
549 if (operand_dim != 1 and operand_dim != out_dim) return error.ShapeMismatch;
550 }
551 for (result_shape) |s| if (s < 0) return error.InvalidDimension;
552 return singletonType(arena, .{
553 .dtype = operand.dtype,
554 .dims = try duplicateDims(arena, result_shape),
555 });
556 }
557
558 pub fn inferReshape(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
559 if (inputs.len != 1) return error.ArityMismatch;
560 const new_shape = try inferAttrI64List(attrs, 0);
561 for (new_shape) |s| if (s < 0) return error.InvalidDimension;
562 const old_count = productI64(inputs[0].dims);
563 const new_count = productI64(new_shape);
564 if (old_count != new_count) return error.ElementCountMismatch;
565 return singletonType(arena, .{
566 .dtype = inputs[0].dtype,
567 .dims = try duplicateDims(arena, new_shape),
568 });
569 }
570
571 fn productI64(dims: []const i64) i64 {
572 var p: i64 = 1;
573 for (dims) |d| p *= d;
574 return p;
575 }
576
577 pub fn inferTranspose(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
578 if (inputs.len != 1) return error.ArityMismatch;
579 const perm = try inferAttrI64List(attrs, 0);
580 const operand = inputs[0];
581 if (perm.len != operand.dims.len) return error.InvalidPermutation;
582 const rank: i64 = @intCast(operand.dims.len);
583 for (perm, 0..) |d, i| {
584 if (d < 0 or d >= rank) return error.InvalidPermutation;
585 for (perm[0..i]) |e| if (e == d) return error.InvalidPermutation;
586 }
587 const out_dims = try arena.alloc(i64, operand.dims.len);
588 for (perm, 0..) |d, i| out_dims[i] = operand.dims[@intCast(d)];
589 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
590 }
591
592 pub fn inferSlice(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
593 if (inputs.len != 1) return error.ArityMismatch;
594 const start = try inferAttrI64List(attrs, 0);
595 const limit = try inferAttrI64List(attrs, 1);
596 const stride = try inferAttrI64List(attrs, 2);
597 const operand = inputs[0];
598 if (start.len != operand.dims.len or
599 limit.len != operand.dims.len or
600 stride.len != operand.dims.len) return error.InvalidDimension;
601 const out_dims = try arena.alloc(i64, operand.dims.len);
602 for (start, limit, stride, operand.dims, 0..) |s, l, st, full, i| {
603 if (st <= 0) return error.InvalidDimension;
604 if (s < 0 or l < s or l > full) return error.InvalidDimension;
605 const span = l - s;
606 out_dims[i] = @divTrunc(span + st - 1, st);
607 }
608 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
609 }
610
611 pub fn inferGather(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
612 if (inputs.len != 2) return error.ArityMismatch;
613 const axis_i64 = try inferAttrI64(attrs, 0);
614 const operand = inputs[0];
615 const indices = inputs[1];
616 if (!isIndexDType(indices.dtype)) return error.DTypeMismatch;
617 const rank: i64 = @intCast(operand.dims.len);
618 if (axis_i64 < 0 or axis_i64 >= rank) return error.InvalidDimension;
619 const axis: usize = @intCast(axis_i64);
620
621 const out_rank = operand.dims.len - 1 + indices.dims.len;
622 const out_dims = try arena.alloc(i64, out_rank);
623 var out_i: usize = 0;
624 for (operand.dims[0..axis]) |dim| {
625 out_dims[out_i] = dim;
626 out_i += 1;
627 }
628 for (indices.dims) |dim| {
629 if (dim < 0) return error.InvalidDimension;
630 out_dims[out_i] = dim;
631 out_i += 1;
632 }
633 for (operand.dims[axis + 1 ..]) |dim| {
634 out_dims[out_i] = dim;
635 out_i += 1;
636 }
637 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
638 }
639
640 pub fn inferScatter(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
641 if (inputs.len != 3) return error.ArityMismatch;
642 const axis_i64 = try inferAttrI64(attrs, 0);
643 const operand = inputs[0];
644 const indices = inputs[1];
645 const updates = inputs[2];
646 if (!isIndexDType(indices.dtype)) return error.DTypeMismatch;
647 if (updates.dtype != operand.dtype) return error.DTypeMismatch;
648 const rank: i64 = @intCast(operand.dims.len);
649 if (axis_i64 < 0 or axis_i64 >= rank) return error.InvalidDimension;
650 const axis: usize = @intCast(axis_i64);
651
652 const expected_rank = operand.dims.len - 1 + indices.dims.len;
653 if (updates.dims.len != expected_rank) return error.RankMismatch;
654
655 var update_axis: usize = 0;
656 for (operand.dims[0..axis]) |dim| {
657 if (updates.dims[update_axis] != dim) return error.ShapeMismatch;
658 update_axis += 1;
659 }
660 for (indices.dims) |dim| {
661 if (dim < 0) return error.InvalidDimension;
662 if (updates.dims[update_axis] != dim) return error.ShapeMismatch;
663 update_axis += 1;
664 }
665 for (operand.dims[axis + 1 ..]) |dim| {
666 if (updates.dims[update_axis] != dim) return error.ShapeMismatch;
667 update_axis += 1;
668 }
669
670 const out_dims = try arena.dupe(i64, operand.dims);
671 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
672 }
673
674 pub fn inferSparseCrossEntropy(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
675 _ = attrs;
676 if (inputs.len != 2) return error.ArityMismatch;
677 const logits = inputs[0];
678 const targets = inputs[1];
679 if (!logits.dtype.isFloat()) return error.DTypeMismatch;
680 if (!isIndexDType(targets.dtype)) return error.DTypeMismatch;
681 if (logits.dims.len != 2 or targets.dims.len != 1) return error.RankMismatch;
682 if (logits.dims[0] < 0 or logits.dims[1] < 0) return error.InvalidDimension;
683 if (targets.dims[0] != logits.dims[0]) return error.ShapeMismatch;
684
685 const out_dims = try arena.dupe(i64, logits.dims[0..1]);
686 return singletonType(arena, .{ .dtype = logits.dtype, .dims = out_dims });
687 }
688
689 pub fn inferPad(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
690 if (inputs.len != 2) return error.ArityMismatch;
691 const low = try inferAttrI64List(attrs, 0);
692 const high = try inferAttrI64List(attrs, 1);
693 const interior = try inferAttrI64List(attrs, 2);
694 const operand = inputs[0];
695 const padding_value = inputs[1];
696 if (padding_value.dtype != operand.dtype) return error.DTypeMismatch;
697 if (padding_value.dims.len != 0) return error.RankMismatch;
698 if (low.len != operand.dims.len or high.len != operand.dims.len or interior.len != operand.dims.len) {
699 return error.InvalidDimension;
700 }
701
702 const out_dims = try arena.alloc(i64, operand.dims.len);
703 for (operand.dims, low, high, interior, 0..) |input_dim, lo, hi, inner, i| {
704 if (input_dim < 0 or inner < 0) return error.InvalidDimension;
705 const gap_count = if (input_dim > 0) input_dim - 1 else 0;
706 const interior_total = std.math.mul(i64, gap_count, inner) catch return error.InvalidDimension;
707 const with_low = std.math.add(i64, input_dim, lo) catch return error.InvalidDimension;
708 const with_high = std.math.add(i64, with_low, hi) catch return error.InvalidDimension;
709 const result_dim = std.math.add(i64, with_high, interior_total) catch return error.InvalidDimension;
710 if (result_dim < 0) return error.InvalidDimension;
711 out_dims[i] = result_dim;
712 }
713 return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });
714 }
715
716 fn isIndexDType(dtype: DType) bool {
717 return dtype.isSignedInt() or dtype.isUnsignedInt();
718 }
719
720 pub fn inferConcatenate(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
721 if (inputs.len == 0) return error.ArityMismatch;
722 const dim = try inferAttrI64(attrs, 0);
723 const first = inputs[0];
724 const rank: i64 = @intCast(first.dims.len);
725 if (dim < 0 or dim >= rank) return error.InvalidDimension;
726 const concat_axis: usize = @intCast(dim);
727
728 var summed = first.dims[concat_axis];
729 for (inputs[1..]) |t| {
730 if (t.dtype != first.dtype) return error.DTypeMismatch;
731 if (t.dims.len != first.dims.len) return error.RankMismatch;
732 for (t.dims, first.dims, 0..) |a, b, i| {
733 if (i == concat_axis) continue;
734 if (a != b) return error.DimMismatch;
735 }
736 summed += t.dims[concat_axis];
737 }
738 const out_dims = try arena.alloc(i64, first.dims.len);
739 @memcpy(out_dims, first.dims);
740 out_dims[concat_axis] = summed;
741 return singletonType(arena, .{ .dtype = first.dtype, .dims = out_dims });
742 }
743
744 pub fn inferScratch(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
745 const words = try inferAttrI64(attrs, 0);
746 if (inputs.len != 0) return error.ArityMismatch;
747 if (words < 1) return error.InvalidDimension;
748 const dims = try arena.alloc(i64, 1);
749 dims[0] = words;
750 return singletonType(arena, .{ .dtype = .u32, .dims = dims });
751 }
752
753 pub fn inferCumsum(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
754 const axis = try inferAttrI64(attrs, 0);
755 if (inputs.len < 1 or inputs.len > 2) return error.ArityMismatch;
756 if (axis < 0 or axis >= @as(i64, @intCast(inputs[0].dims.len))) return error.InvalidDimension;
757 return singletonType(arena, .{
758 .dtype = inputs[0].dtype,
759 .dims = try duplicateDims(arena, inputs[0].dims),
760 });
761 }
762
763 pub fn inferIterate(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
764 const max_iters = try inferAttrI64(attrs, 0);
765 if (max_iters < 1) return error.InvalidDimension;
766 if (inputs.len == 0) return error.ArityMismatch;
767 try verifySingleElementDomain(inputs);
768 const results = try arena.alloc(Type, inputs.len);
769 for (results, inputs) |*result, input| {
770 result.* = .{ .dtype = input.dtype, .dims = try duplicateDims(arena, input.dims) };
771 }
772 return results;
773 }
774
775 pub fn inferIterateYield(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
776 _ = attrs;
777 if (inputs.len < 2) return error.ArityMismatch;
778 if (inputs[0].dtype != .i1) return error.DTypeMismatch;
779 if (typeElementCount(inputs[0]) != 1) {
780 for (inputs[1..]) |carry| {
781 if (typeElementCount(carry) == 1) continue;
782 if (!std.mem.eql(i64, carry.dims, inputs[0].dims)) return error.ShapeMismatch;
783 }
784 }
785 try verifySingleElementDomain(inputs[1..]);
786 return try arena.alloc(Type, 0);
787 }
788
789 fn verifySingleElementDomain(inputs: []const Type) InferError!void {
790 var domain: ?u64 = null;
791 for (inputs) |input| {
792 const count = typeElementCount(input);
793 if (count == 1) continue;
794 if (domain) |existing| {
795 if (count != existing) return error.ShapeMismatch;
796 } else {
797 domain = count;
798 }
799 }
800 }
801
802 fn typeElementCount(input: Type) u64 {
803 var count: u64 = 1;
804 for (input.dims) |dim| {
805 if (dim <= 0) return 0;
806 count *= @intCast(dim);
807 }
808 return count;
809 }
810
811 pub fn inferSelect(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
812 _ = attrs;
813 if (inputs.len != 3) return error.ArityMismatch;
814 const pred = inputs[0];
815 const ot = inputs[1];
816 const of = inputs[2];
817 if (pred.dtype != .i1) return error.DTypeMismatch;
818 if (ot.dtype != of.dtype) return error.DTypeMismatch;
819 if (!std.mem.eql(i64, ot.dims, of.dims)) return error.ShapeMismatch;
820 if (pred.dims.len != 0 and !std.mem.eql(i64, pred.dims, ot.dims)) {
821 return error.ShapeMismatch;
822 }
823 return singletonType(arena, .{
824 .dtype = ot.dtype,
825 .dims = try duplicateDims(arena, ot.dims),
826 });
827 }
828
829 pub fn inferKernelCall(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
830 const contract = try inferAttrKernelCall(attrs, 0);
831 if (contract.target.len == 0) return error.InvalidKernelContract;
832 if (contract.version == 0) return error.InvalidKernelContract;
833 if (contract.results.len == 0) return error.InvalidKernelContract;
834 if (contract.operand_effects.len != inputs.len) return error.InvalidKernelContract;
835 if (contract.result_aliases.len != contract.results.len) return error.InvalidKernelContract;
836 const aliased_operands = try arena.alloc(bool, inputs.len);
837 @memset(aliased_operands, false);
838 for (contract.result_aliases, 0..) |alias, result_index| {
839 const operand_index = alias orelse continue;
840 if (operand_index >= inputs.len) return error.InvalidKernelContract;
841 if (aliased_operands[operand_index]) return error.InvalidKernelContract;
842 aliased_operands[operand_index] = true;
843 if (!contract.operand_effects[operand_index].writes()) return error.InvalidKernelContract;
844 if (!contract.results[result_index].eql(inputs[operand_index])) return error.InvalidKernelContract;
845 }
846 return try duplicateTypes(arena, contract.results);
847 }
848
849 pub fn inferReturn(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {
850 _ = inputs;
851 _ = attrs;
852 return try arena.alloc(Type, 0);
853 }
854
855 const testArena = if (@import("builtin").is_test)
856 struct {
857 fn init(backing: std.mem.Allocator) std.heap.ArenaAllocator {
858 return std.heap.ArenaAllocator.init(backing);
859 }
860 }.init
861 else {};
862
863 test "add infers result type from matching operands" {
864 var a = testArena(testing.allocator);
865 defer a.deinit();
866 const dims = [_]i64{ 4, 8 };
867 const t: Type = .{ .dtype = .f32, .dims = &dims };
868 const out = try inferShape(.add, a.allocator(), &.{ t, t }, &.{});
869 try testing.expectEqual(@as(usize, 1), out.len);
870 try testing.expect(out[0].eql(t));
871 }
872
873 test "add rejects dtype mismatch" {
874 var a = testArena(testing.allocator);
875 defer a.deinit();
876 const dims = [_]i64{ 2, 3 };
877 const t0: Type = .{ .dtype = .f32, .dims = &dims };
878 const t1: Type = .{ .dtype = .i32, .dims = &dims };
879 try testing.expectError(error.DTypeMismatch, inferShape(.add, a.allocator(), &.{ t0, t1 }, &.{}));
880 }
881
882 test "add rejects shape mismatch" {
883 var a = testArena(testing.allocator);
884 defer a.deinit();
885 const d0 = [_]i64{ 2, 3 };
886 const d1 = [_]i64{ 2, 4 };
887 const t0: Type = .{ .dtype = .f32, .dims = &d0 };
888 const t1: Type = .{ .dtype = .f32, .dims = &d1 };
889 try testing.expectError(error.ShapeMismatch, inferShape(.add, a.allocator(), &.{ t0, t1 }, &.{}));
890 }
891
892 test "elementwise binary rejects wrong arity" {
893 var a = testArena(testing.allocator);
894 defer a.deinit();
895 const d = [_]i64{3};
896 const t: Type = .{ .dtype = .f32, .dims = &d };
897 try testing.expectError(error.ArityMismatch, inferShape(.mul, a.allocator(), &.{t}, &.{}));
898 }
899
900 test "compare emits i1 with operand shape" {
901 var a = testArena(testing.allocator);
902 defer a.deinit();
903 const dims = [_]i64{ 4, 8 };
904 const t: Type = .{ .dtype = .f32, .dims = &dims };
905 const attrs = [_]Attribute{.{ .compare_direction = .lt }};
906 const out = try inferShape(.compare, a.allocator(), &.{ t, t }, &attrs);
907 try testing.expectEqual(@as(usize, 1), out.len);
908 try testing.expectEqual(@as(DType, .i1), out[0].dtype);
909 try testing.expectEqualSlices(i64, &dims, out[0].dims);
910 }
911
912 test "compare requires a compare_direction attribute" {
913 var a = testArena(testing.allocator);
914 defer a.deinit();
915 const dims = [_]i64{2};
916 const t: Type = .{ .dtype = .f32, .dims = &dims };
917 try testing.expectError(error.AttributeMissing, inferShape(.compare, a.allocator(), &.{ t, t }, &.{}));
918 }
919
920 test "neg preserves type" {
921 var a = testArena(testing.allocator);
922 defer a.deinit();
923 const dims = [_]i64{6};
924 const t: Type = .{ .dtype = .f32, .dims = &dims };
925 const out = try inferShape(.neg, a.allocator(), &.{t}, &.{});
926 try testing.expect(out[0].eql(t));
927 }
928
929 test "activation preserves type and requires activation kind" {
930 var a = testArena(testing.allocator);
931 defer a.deinit();
932 const dims = [_]i64{8};
933 const t: Type = .{ .dtype = .f32, .dims = &dims };
934 const attrs = [_]Attribute{.{ .activation_kind = .gelu }};
935 const out = try inferShape(.activation, a.allocator(), &.{t}, &attrs);
936 try testing.expect(out[0].eql(t));
937 try testing.expectError(error.AttributeMissing, inferShape(.activation, a.allocator(), &.{t}, &.{}));
938 }
939
940 test "convert changes dtype, keeps shape" {
941 var a = testArena(testing.allocator);
942 defer a.deinit();
943 const dims = [_]i64{ 3, 4 };
944 const t: Type = .{ .dtype = .i32, .dims = &dims };
945 const attrs = [_]Attribute{.{ .dtype = .f32 }};
946 const out = try inferShape(.convert, a.allocator(), &.{t}, &attrs);
947 try testing.expectEqual(@as(DType, .f32), out[0].dtype);
948 try testing.expectEqualSlices(i64, &dims, out[0].dims);
949 }
950
951 test "constant reads dtype + dims from attributes" {
952 var a = testArena(testing.allocator);
953 defer a.deinit();
954 var payload: [16]u8 = @splat(0);
955 const dims = [_]i64{ 2, 2 };
956 const attrs = [_]Attribute{
957 .{ .bytes = &payload },
958 .{ .dtype = .f32 },
959 .{ .i64_list = &dims },
960 };
961 const out = try inferShape(.constant, a.allocator(), &.{}, &attrs);
962 try testing.expectEqual(@as(DType, .f32), out[0].dtype);
963 try testing.expectEqualSlices(i64, &dims, out[0].dims);
964 }
965
966 test "iota rejects out-of-range iota_dimension" {
967 var a = testArena(testing.allocator);
968 defer a.deinit();
969 const dims = [_]i64{ 3, 4 };
970 const attrs = [_]Attribute{
971 .{ .i64 = 2 },
972 .{ .dtype = .i32 },
973 .{ .i64_list = &dims },
974 };
975 try testing.expectError(error.InvalidDimension, inferShape(.iota, a.allocator(), &.{}, &attrs));
976 }
977
978 test "parameter carries dtype + dims" {
979 var a = testArena(testing.allocator);
980 defer a.deinit();
981 const dims = [_]i64{ 4, 8 };
982 const attrs = [_]Attribute{
983 .{ .i64 = 0 },
984 .{ .dtype = .f32 },
985 .{ .i64_list = &dims },
986 };
987 const out = try inferShape(.parameter, a.allocator(), &.{}, &attrs);
988 try testing.expectEqualSlices(i64, &dims, out[0].dims);
989 }
990
991 test "reduce drops reduction dims" {
992 var a = testArena(testing.allocator);
993 defer a.deinit();
994 const in_dims = [_]i64{ 2, 3, 4 };
995 const init_dims = [_]i64{};
996 const red_dims = [_]i64{1};
997 const input: Type = .{ .dtype = .f32, .dims = &in_dims };
998 const init_v: Type = .{ .dtype = .f32, .dims = &init_dims };
999 const attrs = [_]Attribute{
1000 .{ .reducer_kind = .sum },
1001 .{ .i64_list = &red_dims },
1002 };
1003 const out = try inferShape(.reduce, a.allocator(), &.{ input, init_v }, &attrs);
1004 try testing.expectEqualSlices(i64, &[_]i64{ 2, 4 }, out[0].dims);
1005 }
1006
1007 test "reduce rejects non-scalar init value" {
1008 var a = testArena(testing.allocator);
1009 defer a.deinit();
1010 const in_dims = [_]i64{ 2, 3 };
1011 const init_dims = [_]i64{1};
1012 const red_dims = [_]i64{0};
1013 const input: Type = .{ .dtype = .f32, .dims = &in_dims };
1014 const init_v: Type = .{ .dtype = .f32, .dims = &init_dims };
1015 const attrs = [_]Attribute{
1016 .{ .reducer_kind = .sum },
1017 .{ .i64_list = &red_dims },
1018 };
1019 try testing.expectError(error.RankMismatch, inferShape(.reduce, a.allocator(), &.{ input, init_v }, &attrs));
1020 }
1021
1022 test "dot_general matmul: [M,K] x [K,N] -> [M,N]" {
1023 var a = testArena(testing.allocator);
1024 defer a.deinit();
1025 const lhs_d = [_]i64{ 4, 8 };
1026 const rhs_d = [_]i64{ 8, 16 };
1027 const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };
1028 const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };
1029 const empty = [_]i64{};
1030 const lhs_c = [_]i64{1};
1031 const rhs_c = [_]i64{0};
1032 const attrs = [_]Attribute{
1033 .{ .i64_list = &empty },
1034 .{ .i64_list = &empty },
1035 .{ .i64_list = &lhs_c },
1036 .{ .i64_list = &rhs_c },
1037 };
1038 const out = try inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs);
1039 try testing.expectEqualSlices(i64, &[_]i64{ 4, 16 }, out[0].dims);
1040 }
1041
1042 test "dot_general batched: [B,M,K] x [B,K,N] -> [B,M,N]" {
1043 var a = testArena(testing.allocator);
1044 defer a.deinit();
1045 const lhs_d = [_]i64{ 2, 3, 5 };
1046 const rhs_d = [_]i64{ 2, 5, 7 };
1047 const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };
1048 const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };
1049 const b_dims = [_]i64{0};
1050 const lhs_c = [_]i64{2};
1051 const rhs_c = [_]i64{1};
1052 const attrs = [_]Attribute{
1053 .{ .i64_list = &b_dims },
1054 .{ .i64_list = &b_dims },
1055 .{ .i64_list = &lhs_c },
1056 .{ .i64_list = &rhs_c },
1057 };
1058 const out = try inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs);
1059 try testing.expectEqualSlices(i64, &[_]i64{ 2, 3, 7 }, out[0].dims);
1060 }
1061
1062 test "dot_general rejects mismatched contracting dims" {
1063 var a = testArena(testing.allocator);
1064 defer a.deinit();
1065 const lhs_d = [_]i64{ 4, 8 };
1066 const rhs_d = [_]i64{ 9, 16 };
1067 const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };
1068 const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };
1069 const empty = [_]i64{};
1070 const lhs_c = [_]i64{1};
1071 const rhs_c = [_]i64{0};
1072 const attrs = [_]Attribute{
1073 .{ .i64_list = &empty },
1074 .{ .i64_list = &empty },
1075 .{ .i64_list = &lhs_c },
1076 .{ .i64_list = &rhs_c },
1077 };
1078 try testing.expectError(error.ContractionMismatch, inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs));
1079 }
1080
1081 test "einsum infers matrix product output shape" {
1082 var a = testArena(testing.allocator);
1083 defer a.deinit();
1084 const lhs_d = [_]i64{ 4, 8 };
1085 const rhs_d = [_]i64{ 8, 16 };
1086 const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };
1087 const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };
1088 const attrs = [_]Attribute{.{ .einsum = "ik,kj->ij" }};
1089 const out = try inferShape(.einsum, a.allocator(), &.{ lhs, rhs }, &attrs);
1090 try testing.expectEqualSlices(i64, &[_]i64{ 4, 16 }, out[0].dims);
1091 }
1092
1093 test "einsum infers scalar reduction output shape" {
1094 var a = testArena(testing.allocator);
1095 defer a.deinit();
1096 const dims = [_]i64{ 4, 8 };
1097 const typ: Type = .{ .dtype = .f32, .dims = &dims };
1098 const attrs = [_]Attribute{.{ .einsum = "ij->" }};
1099 const out = try inferShape(.einsum, a.allocator(), &.{typ}, &attrs);
1100 try testing.expectEqual(@as(usize, 0), out[0].dims.len);
1101 }
1102
1103 test "einsum rejects inconsistent shared dimensions" {
1104 var a = testArena(testing.allocator);
1105 defer a.deinit();
1106 const lhs_d = [_]i64{ 4, 8 };
1107 const rhs_d = [_]i64{ 9, 16 };
1108 const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };
1109 const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };
1110 const attrs = [_]Attribute{.{ .einsum = "ik,kj->ij" }};
1111 try testing.expectError(error.InvalidEinsum, inferShape(.einsum, a.allocator(), &.{ lhs, rhs }, &attrs));
1112 }
1113
1114 test "broadcast prepends leading dims" {
1115 var a = testArena(testing.allocator);
1116 defer a.deinit();
1117 const in_d = [_]i64{ 3, 4 };
1118 const t: Type = .{ .dtype = .f32, .dims = &in_d };
1119 const sizes = [_]i64{2};
1120 const attrs = [_]Attribute{.{ .i64_list = &sizes }};
1121 const out = try inferShape(.broadcast, a.allocator(), &.{t}, &attrs);
1122 try testing.expectEqualSlices(i64, &[_]i64{ 2, 3, 4 }, out[0].dims);
1123 }
1124
1125 test "broadcast_in_dim matches StableHLO 1x3 -> 2x3x2 example" {
1126 var a = testArena(testing.allocator);
1127 defer a.deinit();
1128 const in_d = [_]i64{ 1, 3 };
1129 const t: Type = .{ .dtype = .i32, .dims = &in_d };
1130 const bcast = [_]i64{ 2, 1 };
1131 const result = [_]i64{ 2, 3, 2 };
1132 const attrs = [_]Attribute{
1133 .{ .i64_list = &bcast },
1134 .{ .i64_list = &result },
1135 };
1136 const out = try inferShape(.broadcast_in_dim, a.allocator(), &.{t}, &attrs);
1137 try testing.expectEqualSlices(i64, &result, out[0].dims);
1138 }
1139
1140 test "broadcast_in_dim rejects incompatible source dim" {
1141 var a = testArena(testing.allocator);
1142 defer a.deinit();
1143 const in_d = [_]i64{ 2, 3 };
1144 const t: Type = .{ .dtype = .i32, .dims = &in_d };
1145 const bcast = [_]i64{ 0, 1 };
1146 const result = [_]i64{ 4, 3 };
1147 const attrs = [_]Attribute{
1148 .{ .i64_list = &bcast },
1149 .{ .i64_list = &result },
1150 };
1151 try testing.expectError(error.ShapeMismatch, inferShape(.broadcast_in_dim, a.allocator(), &.{t}, &attrs));
1152 }
1153
1154 test "reshape changes rank, preserves element count" {
1155 var a = testArena(testing.allocator);
1156 defer a.deinit();
1157 const in_d = [_]i64{ 2, 3 };
1158 const t: Type = .{ .dtype = .i32, .dims = &in_d };
1159 const new = [_]i64{ 3, 2 };
1160 const attrs = [_]Attribute{.{ .i64_list = &new }};
1161 const out = try inferShape(.reshape, a.allocator(), &.{t}, &attrs);
1162 try testing.expectEqualSlices(i64, &new, out[0].dims);
1163 }
1164
1165 test "reshape rejects element-count mismatch" {
1166 var a = testArena(testing.allocator);
1167 defer a.deinit();
1168 const in_d = [_]i64{ 2, 3 };
1169 const t: Type = .{ .dtype = .i32, .dims = &in_d };
1170 const new = [_]i64{ 4, 2 };
1171 const attrs = [_]Attribute{.{ .i64_list = &new }};
1172 try testing.expectError(error.ElementCountMismatch, inferShape(.reshape, a.allocator(), &.{t}, &attrs));
1173 }
1174
1175 test "transpose permutes dims" {
1176 var a = testArena(testing.allocator);
1177 defer a.deinit();
1178 const in_d = [_]i64{ 2, 3, 4 };
1179 const t: Type = .{ .dtype = .f32, .dims = &in_d };
1180 const perm = [_]i64{ 2, 0, 1 };
1181 const attrs = [_]Attribute{.{ .i64_list = &perm }};
1182 const out = try inferShape(.transpose, a.allocator(), &.{t}, &attrs);
1183 try testing.expectEqualSlices(i64, &[_]i64{ 4, 2, 3 }, out[0].dims);
1184 }
1185
1186 test "transpose rejects non-permutation" {
1187 var a = testArena(testing.allocator);
1188 defer a.deinit();
1189 const in_d = [_]i64{ 2, 3, 4 };
1190 const t: Type = .{ .dtype = .f32, .dims = &in_d };
1191 const perm = [_]i64{ 0, 0, 1 };
1192 const attrs = [_]Attribute{.{ .i64_list = &perm }};
1193 try testing.expectError(error.InvalidPermutation, inferShape(.transpose, a.allocator(), &.{t}, &attrs));
1194 }
1195
1196 test "slice computes ceil((limit - start) / stride)" {
1197 var a = testArena(testing.allocator);
1198 defer a.deinit();
1199 const in_d = [_]i64{ 10, 10 };
1200 const t: Type = .{ .dtype = .i32, .dims = &in_d };
1201 const start = [_]i64{ 1, 2 };
1202 const limit = [_]i64{ 9, 8 };
1203 const stride = [_]i64{ 2, 3 };
1204 const attrs = [_]Attribute{
1205 .{ .i64_list = &start },
1206 .{ .i64_list = &limit },
1207 .{ .i64_list = &stride },
1208 };
1209 const out = try inferShape(.slice, a.allocator(), &.{t}, &attrs);
1210 try testing.expectEqualSlices(i64, &[_]i64{ 4, 2 }, out[0].dims);
1211 }
1212
1213 test "gather inserts index shape at the selected axis" {
1214 var a = testArena(testing.allocator);
1215 defer a.deinit();
1216 const operand_dims = [_]i64{ 2, 3, 4 };
1217 const index_dims = [_]i64{ 5, 6 };
1218 const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };
1219 const indices: Type = .{ .dtype = .i32, .dims = &index_dims };
1220 const attrs = [_]Attribute{.{ .i64 = 1 }};
1221 const out = try inferShape(.gather, a.allocator(), &.{ operand, indices }, &attrs);
1222 try testing.expectEqualSlices(i64, &[_]i64{ 2, 5, 6, 4 }, out[0].dims);
1223 }
1224
1225 test "gather rejects invalid axis and non-integer indices" {
1226 var a = testArena(testing.allocator);
1227 defer a.deinit();
1228 const operand_dims = [_]i64{ 2, 3 };
1229 const index_dims = [_]i64{4};
1230 const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };
1231 const indices: Type = .{ .dtype = .i32, .dims = &index_dims };
1232 const bad_indices: Type = .{ .dtype = .f32, .dims = &index_dims };
1233 try testing.expectError(error.InvalidDimension, inferShape(.gather, a.allocator(), &.{ operand, indices }, &.{.{ .i64 = 2 }}));
1234 try testing.expectError(error.DTypeMismatch, inferShape(.gather, a.allocator(), &.{ operand, bad_indices }, &.{.{ .i64 = 0 }}));
1235 }
1236
1237 test "scatter returns operand shape and checks update shape" {
1238 var a = testArena(testing.allocator);
1239 defer a.deinit();
1240 const operand_dims = [_]i64{ 2, 3, 4 };
1241 const index_dims = [_]i64{ 5, 6 };
1242 const update_dims = [_]i64{ 2, 5, 6, 4 };
1243 const bad_update_dims = [_]i64{ 2, 5, 4 };
1244 const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };
1245 const indices: Type = .{ .dtype = .i32, .dims = &index_dims };
1246 const updates: Type = .{ .dtype = .f32, .dims = &update_dims };
1247 const bad_updates: Type = .{ .dtype = .f32, .dims = &bad_update_dims };
1248 const attrs = [_]Attribute{.{ .i64 = 1 }};
1249 const out = try inferShape(.scatter, a.allocator(), &.{ operand, indices, updates }, &attrs);
1250 try testing.expectEqualSlices(i64, &operand_dims, out[0].dims);
1251 try testing.expectError(error.RankMismatch, inferShape(.scatter, a.allocator(), &.{ operand, indices, bad_updates }, &attrs));
1252 }
1253
1254 test "scatter rejects invalid axis and non-integer indices" {
1255 var a = testArena(testing.allocator);
1256 defer a.deinit();
1257 const operand_dims = [_]i64{ 2, 3 };
1258 const index_dims = [_]i64{4};
1259 const update_dims = [_]i64{ 4, 3 };
1260 const operand: Type = .{ .dtype = .i32, .dims = &operand_dims };
1261 const indices: Type = .{ .dtype = .i32, .dims = &index_dims };
1262 const bad_indices: Type = .{ .dtype = .f32, .dims = &index_dims };
1263 const updates: Type = .{ .dtype = .i32, .dims = &update_dims };
1264 try testing.expectError(error.InvalidDimension, inferShape(.scatter, a.allocator(), &.{ operand, indices, updates }, &.{.{ .i64 = 2 }}));
1265 try testing.expectError(error.DTypeMismatch, inferShape(.scatter, a.allocator(), &.{ operand, bad_indices, updates }, &.{.{ .i64 = 0 }}));
1266 }
1267
1268 test "pad infers positive edge and interior padding" {
1269 var a = testArena(testing.allocator);
1270 defer a.deinit();
1271 const in_d = [_]i64{ 2, 3 };
1272 const scalar_d = [_]i64{};
1273 const input: Type = .{ .dtype = .i32, .dims = &in_d };
1274 const padding_value: Type = .{ .dtype = .i32, .dims = &scalar_d };
1275 const low = [_]i64{ 1, 0 };
1276 const high = [_]i64{ 0, 2 };
1277 const interior = [_]i64{ 0, 1 };
1278 const attrs = [_]Attribute{
1279 .{ .i64_list = &low },
1280 .{ .i64_list = &high },
1281 .{ .i64_list = &interior },
1282 };
1283 const out = try inferShape(.pad, a.allocator(), &.{ input, padding_value }, &attrs);
1284 try testing.expectEqualSlices(i64, &[_]i64{ 3, 7 }, out[0].dims);
1285 }
1286
1287 test "pad allows negative edge padding when result shape remains valid" {
1288 var a = testArena(testing.allocator);
1289 defer a.deinit();
1290 const in_d = [_]i64{4};
1291 const scalar_d = [_]i64{};
1292 const input: Type = .{ .dtype = .f32, .dims = &in_d };
1293 const padding_value: Type = .{ .dtype = .f32, .dims = &scalar_d };
1294 const low = [_]i64{-1};
1295 const high = [_]i64{2};
1296 const interior = [_]i64{0};
1297 const attrs = [_]Attribute{
1298 .{ .i64_list = &low },
1299 .{ .i64_list = &high },
1300 .{ .i64_list = &interior },
1301 };
1302 const out = try inferShape(.pad, a.allocator(), &.{ input, padding_value }, &attrs);
1303 try testing.expectEqualSlices(i64, &[_]i64{5}, out[0].dims);
1304 }
1305
1306 test "pad rejects non-scalar padding value and invalid interior padding" {
1307 var a = testArena(testing.allocator);
1308 defer a.deinit();
1309 const in_d = [_]i64{ 2, 3 };
1310 const scalar_d = [_]i64{};
1311 const vector_d = [_]i64{1};
1312 const input: Type = .{ .dtype = .i32, .dims = &in_d };
1313 const scalar_padding: Type = .{ .dtype = .i32, .dims = &scalar_d };
1314 const vector_padding: Type = .{ .dtype = .i32, .dims = &vector_d };
1315 const low = [_]i64{ 0, 0 };
1316 const high = [_]i64{ 0, 0 };
1317 const interior = [_]i64{ 0, -1 };
1318 const attrs = [_]Attribute{
1319 .{ .i64_list = &low },
1320 .{ .i64_list = &high },
1321 .{ .i64_list = &low },
1322 };
1323 try testing.expectError(error.RankMismatch, inferShape(.pad, a.allocator(), &.{ input, vector_padding }, &attrs));
1324
1325 const bad_attrs = [_]Attribute{
1326 .{ .i64_list = &low },
1327 .{ .i64_list = &high },
1328 .{ .i64_list = &interior },
1329 };
1330 try testing.expectError(error.InvalidDimension, inferShape(.pad, a.allocator(), &.{ input, scalar_padding }, &bad_attrs));
1331 }
1332
1333 test "concatenate sums along the concat dim" {
1334 var a = testArena(testing.allocator);
1335 defer a.deinit();
1336 const d0 = [_]i64{ 3, 2 };
1337 const d1 = [_]i64{ 1, 2 };
1338 const t0: Type = .{ .dtype = .i64, .dims = &d0 };
1339 const t1: Type = .{ .dtype = .i64, .dims = &d1 };
1340 const attrs = [_]Attribute{.{ .i64 = 0 }};
1341 const out = try inferShape(.concatenate, a.allocator(), &.{ t0, t1 }, &attrs);
1342 try testing.expectEqualSlices(i64, &[_]i64{ 4, 2 }, out[0].dims);
1343 }
1344
1345 test "concatenate rejects non-concat-dim mismatch" {
1346 var a = testArena(testing.allocator);
1347 defer a.deinit();
1348 const d0 = [_]i64{ 3, 2 };
1349 const d1 = [_]i64{ 1, 3 };
1350 const t0: Type = .{ .dtype = .i64, .dims = &d0 };
1351 const t1: Type = .{ .dtype = .i64, .dims = &d1 };
1352 const attrs = [_]Attribute{.{ .i64 = 0 }};
1353 try testing.expectError(error.DimMismatch, inferShape(.concatenate, a.allocator(), &.{ t0, t1 }, &attrs));
1354 }
1355
1356 test "select requires i1 predicate and matching branches" {
1357 var a = testArena(testing.allocator);
1358 defer a.deinit();
1359 const d = [_]i64{ 2, 2 };
1360 const pred: Type = .{ .dtype = .i1, .dims = &d };
1361 const on_true: Type = .{ .dtype = .i32, .dims = &d };
1362 const on_false: Type = .{ .dtype = .i32, .dims = &d };
1363 const out = try inferShape(.select, a.allocator(), &.{ pred, on_true, on_false }, &.{});
1364 try testing.expect(out[0].eql(on_true));
1365
1366 const bad_pred: Type = .{ .dtype = .i32, .dims = &d };
1367 try testing.expectError(error.DTypeMismatch, inferShape(.select, a.allocator(), &.{ bad_pred, on_true, on_false }, &.{}));
1368 }
1369
1370 test "kernel_call returns explicit contract result types" {
1371 var a = testArena(testing.allocator);
1372 defer a.deinit();
1373 const input_dims = [_]i64{ 4, 8 };
1374 const result_dims = [_]i64{ 4, 8 };
1375 const input: Type = .{ .dtype = .f32, .dims = &input_dims };
1376 const result: Type = .{ .dtype = .f32, .dims = &result_dims };
1377 const contract = KernelCallContract{
1378 .target = "scale_f32",
1379 .version = 1,
1380 .has_side_effects = false,
1381 .operand_effects = &.{.none},
1382 .result_aliases = &.{null},
1383 .results = &.{result},
1384 };
1385 const out = try inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = contract }});
1386 try testing.expectEqual(@as(usize, 1), out.len);
1387 try testing.expect(out[0].eql(result));
1388 }
1389
1390 test "kernel_call accepts explicit result aliases for writable operands" {
1391 var a = testArena(testing.allocator);
1392 defer a.deinit();
1393 const dims = [_]i64{4};
1394 const input: Type = .{ .dtype = .f32, .dims = &dims };
1395 const contract = KernelCallContract{
1396 .target = "update_f32",
1397 .version = 1,
1398 .has_side_effects = false,
1399 .operand_effects = &.{.read_write},
1400 .result_aliases = &.{0},
1401 .results = &.{input},
1402 };
1403 const out = try inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = contract }});
1404 try testing.expectEqual(@as(usize, 1), out.len);
1405 try testing.expect(out[0].eql(input));
1406 }
1407
1408 test "kernel_call rejects malformed contracts" {
1409 var a = testArena(testing.allocator);
1410 defer a.deinit();
1411 const dims = [_]i64{4};
1412 const input: Type = .{ .dtype = .f32, .dims = &dims };
1413 const result: Type = .{ .dtype = .f32, .dims = &dims };
1414 try testing.expectError(
1415 error.InvalidKernelContract,
1416 inferShape(.kernel_call, a.allocator(), &.{}, &.{.{ .kernel_call = .{
1417 .target = "",
1418 .version = 1,
1419 .has_side_effects = false,
1420 .operand_effects = &.{},
1421 .result_aliases = &.{null},
1422 .results = &.{result},
1423 } }}),
1424 );
1425 try testing.expectError(
1426 error.InvalidKernelContract,
1427 inferShape(.kernel_call, a.allocator(), &.{}, &.{.{ .kernel_call = .{
1428 .target = "scale_f32",
1429 .version = 0,
1430 .has_side_effects = false,
1431 .operand_effects = &.{},
1432 .result_aliases = &.{null},
1433 .results = &.{result},
1434 } }}),
1435 );
1436 try testing.expectError(
1437 error.InvalidKernelContract,
1438 inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{
1439 .target = "scale_f32",
1440 .version = 1,
1441 .has_side_effects = false,
1442 .operand_effects = &.{},
1443 .result_aliases = &.{null},
1444 .results = &.{result},
1445 } }}),
1446 );
1447 try testing.expectError(
1448 error.InvalidKernelContract,
1449 inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{
1450 .target = "update_f32",
1451 .version = 1,
1452 .has_side_effects = false,
1453 .operand_effects = &.{.read},
1454 .result_aliases = &.{0},
1455 .results = &.{result},
1456 } }}),
1457 );
1458 try testing.expectError(
1459 error.InvalidKernelContract,
1460 inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{
1461 .target = "update_f32",
1462 .version = 1,
1463 .has_side_effects = false,
1464 .operand_effects = &.{.read_write},
1465 .result_aliases = &.{ 0, 0 },
1466 .results = &.{ result, result },
1467 } }}),
1468 );
1469 }
1470
1471 test "return has zero outputs" {
1472 var a = testArena(testing.allocator);
1473 defer a.deinit();
1474 const d = [_]i64{3};
1475 const t: Type = .{ .dtype = .f32, .dims = &d };
1476 const out = try inferShape(.@"return", a.allocator(), &.{t}, &.{});
1477 try testing.expectEqual(@as(usize, 0), out.len);
1478 }