lib/accy/src/choir/semantic.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const choir_abi = @import("choir_abi");
   3 const choir = @import("choir");
   4 const dialect_mod = @import("dialect.zig");
   5 const semantics_mod = @import("semantics.zig");
   6 
   7 const ir = choir.ir;
   8 const dialects = choir.dialects;
   9 const AccyDialect = dialect_mod.AccyDialect;
  10 
  11 pub const product_name = "accy.semantic";
  12 pub const ActivationKind = semantics_mod.ActivationKind;
  13 pub const CompareDirection = semantics_mod.CompareDirection;
  14 pub const KernelOperandEffect = semantics_mod.KernelOperandEffect;
  15 
  16 pub const KernelCallOptions = struct {
  17     target: []const u8,
  18     version: u32 = 1,
  19     has_side_effects: bool = false,
  20     operand_effects: []const KernelOperandEffect,
  21     result_aliases: []const ?usize,
  22 };
  23 
  24 pub const SemanticModule = struct {
  25     allocator: std.mem.Allocator,
  26     ctx: *ir.Context,
  27     owns_ctx: bool,
  28     choir_module: *ir.Operation,
  29     live: bool = true,
  30 
  31     pub fn context(self: *SemanticModule) *ir.Context {
  32         return self.ctx;
  33     }
  34 
  35     pub fn deinit(self: *SemanticModule) void {
  36         if (self.live) {
  37             self.choir_module.erase();
  38             if (self.owns_ctx) {
  39                 self.ctx.deinit(self.allocator);
  40                 self.allocator.destroy(self.ctx);
  41             }
  42             self.live = false;
  43         }
  44         const allocator = self.allocator;
  45         allocator.destroy(self);
  46     }
  47 
  48     pub fn verify(self: *SemanticModule) !void {
  49         try ir.verifyOperation(self.choir_module, ir.verify.default_options);
  50     }
  51 
  52     pub fn fingerprint(self: *SemanticModule, allocator: std.mem.Allocator) !choir.product.incremental.Fingerprint {
  53         return try choir.operationFingerprint(allocator, self.choir_module);
  54     }
  55 
  56     pub fn capture(
  57         self: *SemanticModule,
  58         builder: *choir.product.revision.Builder,
  59         comptime configuration: choir.product.operation.Configuration,
  60     ) !void {
  61         try choir.product.operation.capture(
  62             builder,
  63             &.{.{ .operation = self.choir_module }},
  64             &@import("root.zig").publication.irRecord(.semantic),
  65             &.{},
  66             configuration,
  67         );
  68     }
  69 };
  70 
  71 pub const Builder = struct {
  72     module: ?*SemanticModule,
  73     location: ir.Location,
  74 
  75     pub const ContextLimits = ir.Context.Limits;
  76 
  77     pub fn init(
  78         allocator: std.mem.Allocator,
  79         context_limits: ContextLimits,
  80     ) !Builder {
  81         return try initAt(allocator, context_limits, ir.Location.getUnknown());
  82     }
  83 
  84     pub fn initAt(
  85         allocator: std.mem.Allocator,
  86         context_limits: ContextLimits,
  87         location: ir.Location,
  88     ) !Builder {
  89         const ctx = try allocator.create(ir.Context);
  90         errdefer allocator.destroy(ctx);
  91         ctx.* = try buildSemanticContext(allocator, context_limits);
  92         errdefer ctx.deinit(allocator);
  93         return try initIn(allocator, ctx, true, location);
  94     }
  95 
  96     pub fn initBorrowing(allocator: std.mem.Allocator, ctx: *ir.Context) !Builder {
  97         return try initBorrowingAt(allocator, ctx, ir.Location.getUnknown());
  98     }
  99 
 100     pub fn initBorrowingAt(allocator: std.mem.Allocator, ctx: *ir.Context, location: ir.Location) !Builder {
 101         return try initIn(allocator, ctx, false, location);
 102     }
 103 
 104     fn initIn(allocator: std.mem.Allocator, ctx: *ir.Context, owns_ctx: bool, location: ir.Location) !Builder {
 105         const module = try allocator.create(SemanticModule);
 106         errdefer allocator.destroy(module);
 107 
 108         module.* = .{
 109             .allocator = allocator,
 110             .ctx = ctx,
 111             .owns_ctx = owns_ctx,
 112             .choir_module = undefined,
 113         };
 114 
 115         const choir_module = try dialects.BuiltinDialect.ModuleOp.create(module.ctx, location);
 116         module.choir_module = choir_module.op;
 117         return .{ .module = module, .location = location };
 118     }
 119 
 120     pub fn deinit(self: *Builder) void {
 121         if (self.module) |module| {
 122             module.deinit();
 123             self.module = null;
 124         }
 125     }
 126 
 127     pub fn tensor(self: *Builder, dtype: choir_abi.DType, dims: []const i64) !ir.Type {
 128         return try dialect_mod.accyTensorType(self.context(), dtype, dims);
 129     }
 130 
 131     pub fn beginFunction(
 132         self: *Builder,
 133         name: []const u8,
 134         param_types: []const ir.Type,
 135         result_types: []const ir.Type,
 136     ) !FunctionBuilder {
 137         return try self.beginFunctionAt(name, param_types, result_types, self.location);
 138     }
 139 
 140     pub fn beginFunctionAt(
 141         self: *Builder,
 142         name: []const u8,
 143         param_types: []const ir.Type,
 144         result_types: []const ir.Type,
 145         location: ir.Location,
 146     ) !FunctionBuilder {
 147         const module = self.module orelse return error.BuilderFinished;
 148         const func = try dialects.FuncDialect.FuncOp.create(
 149             module.ctx,
 150             location,
 151             name,
 152             param_types,
 153             result_types,
 154         );
 155         errdefer if (func.op.getBlock() == null) func.op.erase();
 156         try bodyBlock(module.choir_module).addOperation(func.op);
 157         return .{
 158             .ctx = module.ctx,
 159             .entry = func.getEntryBlock(),
 160             .result_types = result_types,
 161             .location = location,
 162         };
 163     }
 164 
 165     pub fn finish(self: *Builder) !*SemanticModule {
 166         const module = self.module orelse return error.BuilderFinished;
 167         try module.verify();
 168         self.module = null;
 169         return module;
 170     }
 171 
 172     fn context(self: *Builder) *ir.Context {
 173         return (self.module orelse unreachable).ctx;
 174     }
 175 };
 176 
 177 pub const IterateBuilder = struct {
 178     op: AccyDialect.IterateOp,
 179     body: FunctionBuilder,
 180 
 181     pub fn carry(self: *IterateBuilder, idx: usize) *ir.Value {
 182         return self.body.parameter(idx);
 183     }
 184 
 185     pub fn inner(self: *IterateBuilder) *FunctionBuilder {
 186         return &self.body;
 187     }
 188 
 189     pub fn yield_(self: *IterateBuilder, predicate: *ir.Value, carries: []const *ir.Value) !void {
 190         if (self.body.terminated) return error.AlreadyTerminated;
 191         errdefer self.op.op.erase();
 192         const yield_op = try AccyDialect.IterateYieldOp.create(self.body.ctx, self.body.location, predicate, carries);
 193         try self.body.insert(yield_op.op);
 194         self.body.terminated = true;
 195         try verifyInsertedAccyOp(self.op.op);
 196     }
 197 
 198     pub fn result(self: *IterateBuilder, idx: usize) *ir.Value {
 199         return self.op.getResult(idx).?;
 200     }
 201 };
 202 
 203 pub const FunctionBuilder = struct {
 204     ctx: *ir.Context,
 205     entry: *ir.Block,
 206     result_types: []const ir.Type,
 207     location: ir.Location,
 208     terminated: bool = false,
 209 
 210     pub fn setLocation(self: *FunctionBuilder, location: ir.Location) void {
 211         self.location = location;
 212     }
 213 
 214     pub fn parameter(self: *FunctionBuilder, idx: usize) *ir.Value {
 215         return self.entry.arguments.items[idx];
 216     }
 217 
 218     pub fn constant(self: *FunctionBuilder, result_type: ir.Type, bytes: []const u8) !*ir.Value {
 219         const op = try AccyDialect.ConstantOp.create(self.ctx, self.location, bytes, result_type);
 220         try self.insert(op.op);
 221         return op.getResult();
 222     }
 223 
 224     pub fn iota(self: *FunctionBuilder, result_type: ir.Type, axis: u32) !*ir.Value {
 225         const op = try AccyDialect.IotaOp.create(self.ctx, self.location, result_type, @intCast(axis));
 226         try self.insert(op.op);
 227         return op.getResult();
 228     }
 229 
 230     pub fn add(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 231         const op = try AccyDialect.AddOp.create(self.ctx, self.location, lhs, rhs);
 232         try self.insert(op.op);
 233         return op.getResult();
 234     }
 235 
 236     pub fn sub(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 237         const op = try AccyDialect.SubOp.create(self.ctx, self.location, lhs, rhs);
 238         try self.insert(op.op);
 239         return op.getResult();
 240     }
 241 
 242     pub fn mul(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 243         const op = try AccyDialect.MulOp.create(self.ctx, self.location, lhs, rhs);
 244         try self.insert(op.op);
 245         return op.getResult();
 246     }
 247 
 248     pub fn div(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 249         const op = try AccyDialect.DivOp.create(self.ctx, self.location, lhs, rhs);
 250         try self.insert(op.op);
 251         return op.getResult();
 252     }
 253 
 254     pub fn max(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 255         const op = try AccyDialect.MaxOp.create(self.ctx, self.location, lhs, rhs);
 256         try self.insert(op.op);
 257         return op.getResult();
 258     }
 259 
 260     pub fn min(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 261         const op = try AccyDialect.MinOp.create(self.ctx, self.location, lhs, rhs);
 262         try self.insert(op.op);
 263         return op.getResult();
 264     }
 265 
 266     pub fn neg(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 267         const op = try AccyDialect.NegOp.create(self.ctx, self.location, input);
 268         try self.insert(op.op);
 269         return op.getResult();
 270     }
 271 
 272     pub fn abs(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 273         const op = try AccyDialect.AbsOp.create(self.ctx, self.location, input);
 274         try self.insert(op.op);
 275         return op.getResult();
 276     }
 277 
 278     pub fn exp(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 279         const op = try AccyDialect.ExpOp.create(self.ctx, self.location, input);
 280         try self.insert(op.op);
 281         return op.getResult();
 282     }
 283 
 284     pub fn log(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 285         const op = try AccyDialect.LogOp.create(self.ctx, self.location, input);
 286         try self.insert(op.op);
 287         return op.getResult();
 288     }
 289 
 290     pub fn sqrt(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 291         const op = try AccyDialect.SqrtOp.create(self.ctx, self.location, input);
 292         try self.insert(op.op);
 293         return op.getResult();
 294     }
 295 
 296     pub fn tanh(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 297         const op = try AccyDialect.TanhOp.create(self.ctx, self.location, input);
 298         try self.insert(op.op);
 299         return op.getResult();
 300     }
 301 
 302     pub fn activation(self: *FunctionBuilder, input: *ir.Value, kind: ActivationKind) !*ir.Value {
 303         const op = try AccyDialect.ActivationOp.create(self.ctx, self.location, input, input.type, kind);
 304         try self.insert(op.op);
 305         return op.getResult();
 306     }
 307 
 308     pub fn compare(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value, result_type: ir.Type, direction: CompareDirection) !*ir.Value {
 309         const op = try AccyDialect.CompareOp.create(self.ctx, self.location, lhs, rhs, result_type, @tagName(direction));
 310         try self.insert(op.op);
 311         return op.getResult();
 312     }
 313 
 314     pub fn convert(self: *FunctionBuilder, input: *ir.Value, result_type: ir.Type, target_dtype: choir_abi.DType) !*ir.Value {
 315         const op = try AccyDialect.ConvertOp.create(self.ctx, self.location, input, result_type, @tagName(target_dtype));
 316         try self.insert(op.op);
 317         return op.getResult();
 318     }
 319 
 320     pub fn select(self: *FunctionBuilder, condition: *ir.Value, on_true: *ir.Value, on_false: *ir.Value) !*ir.Value {
 321         const op = try AccyDialect.SelectOp.create(self.ctx, self.location, condition, on_true, on_false);
 322         try self.insert(op.op);
 323         return op.getResult();
 324     }
 325 
 326     pub fn sin(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 327         const op = try AccyDialect.SinOp.create(self.ctx, self.location, input);
 328         try self.insert(op.op);
 329         return op.getResult();
 330     }
 331 
 332     pub fn cos(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 333         const op = try AccyDialect.CosOp.create(self.ctx, self.location, input);
 334         try self.insert(op.op);
 335         return op.getResult();
 336     }
 337 
 338     pub fn tan(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 339         const op = try AccyDialect.TanOp.create(self.ctx, self.location, input);
 340         try self.insert(op.op);
 341         return op.getResult();
 342     }
 343 
 344     pub fn floor(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 345         const op = try AccyDialect.FloorOp.create(self.ctx, self.location, input);
 346         try self.insert(op.op);
 347         return op.getResult();
 348     }
 349 
 350     pub fn round(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 351         const op = try AccyDialect.RoundOp.create(self.ctx, self.location, input);
 352         try self.insert(op.op);
 353         return op.getResult();
 354     }
 355 
 356     pub fn trunc(self: *FunctionBuilder, input: *ir.Value) !*ir.Value {
 357         const op = try AccyDialect.TruncOp.create(self.ctx, self.location, input);
 358         try self.insert(op.op);
 359         return op.getResult();
 360     }
 361 
 362     pub fn pow(self: *FunctionBuilder, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
 363         const op = try AccyDialect.PowOp.create(self.ctx, self.location, lhs, rhs);
 364         try self.insert(op.op);
 365         return op.getResult();
 366     }
 367 
 368     pub fn atan2(self: *FunctionBuilder, y: *ir.Value, x: *ir.Value) !*ir.Value {
 369         const op = try AccyDialect.Atan2Op.create(self.ctx, self.location, y, x);
 370         try self.insert(op.op);
 371         return op.getResult();
 372     }
 373 
 374     pub fn broadcast(self: *FunctionBuilder, input: *ir.Value, result_type: ir.Type, sizes: []const i64) !*ir.Value {
 375         const op = try AccyDialect.BroadcastOp.create(self.ctx, self.location, input, result_type, sizes);
 376         try self.insert(op.op);
 377         return op.getResult();
 378     }
 379 
 380     pub fn broadcastInDim(
 381         self: *FunctionBuilder,
 382         input: *ir.Value,
 383         result_type: ir.Type,
 384         result_shape: []const i64,
 385         broadcast_dims: []const i64,
 386     ) !*ir.Value {
 387         const op = try AccyDialect.BroadcastInDimOp.create(
 388             self.ctx,
 389             self.location,
 390             input,
 391             result_type,
 392             broadcast_dims,
 393             result_shape,
 394         );
 395         try self.insert(op.op);
 396         return op.getResult();
 397     }
 398 
 399     pub fn reshape(self: *FunctionBuilder, input: *ir.Value, result_type: ir.Type, new_shape: []const i64) !*ir.Value {
 400         const op = try AccyDialect.ReshapeOp.create(self.ctx, self.location, input, result_type, new_shape);
 401         try self.insert(op.op);
 402         return op.getResult();
 403     }
 404 
 405     pub fn transpose(self: *FunctionBuilder, input: *ir.Value, result_type: ir.Type, permutation: []const i64) !*ir.Value {
 406         const op = try AccyDialect.TransposeOp.create(self.ctx, self.location, input, result_type, permutation);
 407         try self.insert(op.op);
 408         return op.getResult();
 409     }
 410 
 411     pub fn slice(
 412         self: *FunctionBuilder,
 413         input: *ir.Value,
 414         result_type: ir.Type,
 415         starts: []const i64,
 416         limits: []const i64,
 417         strides: []const i64,
 418     ) !*ir.Value {
 419         const op = try AccyDialect.SliceOp.create(
 420             self.ctx,
 421             self.location,
 422             input,
 423             result_type,
 424             starts,
 425             limits,
 426             strides,
 427         );
 428         try self.insert(op.op);
 429         return op.getResult();
 430     }
 431 
 432     pub fn concatenate(
 433         self: *FunctionBuilder,
 434         operands: []const *ir.Value,
 435         result_type: ir.Type,
 436         dimension: i64,
 437     ) !*ir.Value {
 438         const op = try AccyDialect.ConcatenateOp.create(self.ctx, self.location, operands, result_type, dimension);
 439         try self.insert(op.op);
 440         return op.getResult();
 441     }
 442 
 443     pub fn kernelCall(
 444         self: *FunctionBuilder,
 445         operands: []const *ir.Value,
 446         result_types: []const ir.Type,
 447         options: KernelCallOptions,
 448     ) !AccyDialect.KernelCallOp {
 449         const op = try AccyDialect.KernelCallOp.create(
 450             self.ctx,
 451             self.location,
 452             operands,
 453             result_types,
 454             options.target,
 455             options.version,
 456             options.has_side_effects,
 457             options.operand_effects,
 458             options.result_aliases,
 459         );
 460         try self.insert(op.op);
 461         return op;
 462     }
 463 
 464     pub fn dotGeneral(
 465         self: *FunctionBuilder,
 466         lhs: *ir.Value,
 467         rhs: *ir.Value,
 468         result_type: ir.Type,
 469         contracting_lhs: []const i64,
 470         contracting_rhs: []const i64,
 471         batching_lhs: []const i64,
 472         batching_rhs: []const i64,
 473     ) !*ir.Value {
 474         const op = try AccyDialect.DotGeneralOp.create(
 475             self.ctx,
 476             self.location,
 477             lhs,
 478             rhs,
 479             result_type,
 480             batching_lhs,
 481             batching_rhs,
 482             contracting_lhs,
 483             contracting_rhs,
 484         );
 485         try self.insert(op.op);
 486         return op.getResult();
 487     }
 488 
 489     pub fn einsum(
 490         self: *FunctionBuilder,
 491         operands: []const *ir.Value,
 492         result_type: ir.Type,
 493         equation: []const u8,
 494     ) !*ir.Value {
 495         const op = try AccyDialect.EinsumOp.create(
 496             self.ctx,
 497             self.location,
 498             operands,
 499             result_type,
 500             equation,
 501         );
 502         try self.insert(op.op);
 503         return op.getResult();
 504     }
 505 
 506     pub fn reduce(
 507         self: *FunctionBuilder,
 508         input: *ir.Value,
 509         init: *ir.Value,
 510         result_type: ir.Type,
 511         reducer_kind: []const u8,
 512         dimensions: []const i64,
 513     ) !*ir.Value {
 514         const op = try AccyDialect.ReduceOp.create(
 515             self.ctx,
 516             self.location,
 517             input,
 518             init,
 519             result_type,
 520             reducer_kind,
 521             dimensions,
 522         );
 523         try self.insert(op.op);
 524         return op.getResult();
 525     }
 526 
 527     pub fn gather(
 528         self: *FunctionBuilder,
 529         input: *ir.Value,
 530         indices: *ir.Value,
 531         result_type: ir.Type,
 532         axis: i64,
 533     ) !*ir.Value {
 534         const op = try AccyDialect.GatherOp.create(self.ctx, self.location, input, indices, result_type, axis);
 535         try self.insert(op.op);
 536         return op.getResult();
 537     }
 538 
 539     pub fn scatter(
 540         self: *FunctionBuilder,
 541         input: *ir.Value,
 542         indices: *ir.Value,
 543         updates: *ir.Value,
 544         result_type: ir.Type,
 545         axis: i64,
 546     ) !*ir.Value {
 547         const op = try AccyDialect.ScatterOp.create(self.ctx, self.location, input, indices, updates, result_type, axis);
 548         try self.insert(op.op);
 549         return op.getResult();
 550     }
 551 
 552     pub fn scatterAdd(
 553         self: *FunctionBuilder,
 554         input: *ir.Value,
 555         indices: *ir.Value,
 556         updates: *ir.Value,
 557         result_type: ir.Type,
 558         axis: i64,
 559     ) !*ir.Value {
 560         const op = try AccyDialect.ScatterAddOp.create(self.ctx, self.location, input, indices, updates, result_type, axis);
 561         try self.insert(op.op);
 562         return op.getResult();
 563     }
 564 
 565     pub fn sparseCrossEntropy(
 566         self: *FunctionBuilder,
 567         logits: *ir.Value,
 568         targets: *ir.Value,
 569         result_type: ir.Type,
 570     ) !*ir.Value {
 571         const op = try AccyDialect.SparseCrossEntropyOp.create(self.ctx, self.location, logits, targets, result_type);
 572         try self.insert(op.op);
 573         return op.getResult();
 574     }
 575 
 576     pub fn pad(
 577         self: *FunctionBuilder,
 578         input: *ir.Value,
 579         padding_value: *ir.Value,
 580         result_type: ir.Type,
 581         edge_low: []const i64,
 582         edge_high: []const i64,
 583         interior: []const i64,
 584     ) !*ir.Value {
 585         const op = try AccyDialect.PadOp.create(
 586             self.ctx,
 587             self.location,
 588             input,
 589             padding_value,
 590             result_type,
 591             edge_low,
 592             edge_high,
 593             interior,
 594         );
 595         try self.insert(op.op);
 596         return op.getResult();
 597     }
 598 
 599     pub fn cumsum(self: *FunctionBuilder, input: *ir.Value, result_type: ir.Type, axis: i64) !*ir.Value {
 600         const op = try AccyDialect.CumsumOp.create(self.ctx, self.location, input, result_type, axis);
 601         try self.insert(op.op);
 602         return op.getResult();
 603     }
 604 
 605     pub fn beginIterate(self: *FunctionBuilder, carries: []const *ir.Value, max_iters: i64) !IterateBuilder {
 606         if (self.terminated) return error.AlreadyTerminated;
 607         const op = try AccyDialect.IterateOp.create(self.ctx, self.location, carries, max_iters);
 608         errdefer if (op.op.getBlock() == null) op.op.erase();
 609         const block = op.bodyBlock() orelse return error.BuilderFinished;
 610         try self.entry.addOperation(op.op);
 611         return .{
 612             .op = op,
 613             .body = .{
 614                 .ctx = self.ctx,
 615                 .entry = block,
 616                 .result_types = &.{},
 617                 .location = self.location,
 618             },
 619         };
 620     }
 621 
 622     pub fn return_(self: *FunctionBuilder, values: []const *ir.Value) !void {
 623         if (self.terminated) return error.AlreadyTerminated;
 624         if (values.len != self.result_types.len) return error.ResultCountMismatch;
 625         for (values, self.result_types) |value, expected| {
 626             if (!value.type.eql(expected)) return error.ResultTypeMismatch;
 627         }
 628         const op = try dialects.FuncDialect.ReturnOp.create(self.ctx, self.location, values);
 629         try self.insert(op.op);
 630         self.terminated = true;
 631     }
 632 
 633     pub fn finish(self: *FunctionBuilder) !void {
 634         if (!self.terminated) return error.MissingTerminator;
 635     }
 636 
 637     fn insert(self: *FunctionBuilder, op: *ir.Operation) !void {
 638         std.debug.assert(op.getBlock() == null);
 639         errdefer if (op.getBlock() == null) op.erase();
 640         if (self.terminated) return error.AlreadyTerminated;
 641         try verifyInsertedAccyOp(op);
 642         try self.entry.addOperation(op);
 643     }
 644 };
 645 
 646 fn verifyInsertedAccyOp(op: *ir.Operation) !void {
 647     if (std.mem.startsWith(u8, op.name.name, "accy.")) {
 648         try ir.verifyOperation(op, .{ .recursive = false });
 649     }
 650 }
 651 
 652 pub fn buildSemanticContext(
 653     allocator: std.mem.Allocator,
 654     context_limits: ir.Context.Limits,
 655 ) !ir.Context {
 656     return try choir.compiler.initContext(allocator, .{
 657         .context_limits = context_limits,
 658         .packages = &.{dialect_mod.accy_package_extension},
 659         .preload_dialects = &.{
 660             "builtin",
 661             "arith",
 662             "memref",
 663             "scf",
 664             "func",
 665             "accy",
 666         },
 667     });
 668 }
 669 
 670 fn bodyBlock(module: *ir.Operation) *ir.Block {
 671     return module.getRegion(0).?.getEntryBlock().?;
 672 }
 673 
 674 const SemanticResourceCounts = struct {
 675     operations: usize,
 676 
 677     fn capture(ctx: *const ir.Context) SemanticResourceCounts {
 678         return .{
 679             .operations = ctx.operationCount(),
 680         };
 681     }
 682 
 683     fn expectEqual(self: SemanticResourceCounts, ctx: *const ir.Context) !void {
 684         try std.testing.expectEqual(self.operations, ctx.operationCount());
 685     }
 686 };
 687 
 688 test "semantic builder emits verified Accy Choir module" {
 689     const allocator = std.testing.allocator;
 690 
 691     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 692     defer builder.deinit();
 693 
 694     const ty = try builder.tensor(.f32, &.{4});
 695     var function = try builder.beginFunction("semantic_add", &.{ ty, ty }, &.{ty});
 696     const sum = try function.add(function.parameter(0), function.parameter(1));
 697     try function.return_(&.{sum});
 698     try function.finish();
 699 
 700     const module = try builder.finish();
 701     defer module.deinit();
 702 
 703     try std.testing.expectEqualStrings(product_name, "accy.semantic");
 704     try std.testing.expect(module.context().isDialectLoaded("accy"));
 705     try std.testing.expect(module.context().isDialectLoaded("func"));
 706     try module.verify();
 707     try std.testing.expect(try module.fingerprint(allocator) != 0);
 708 }
 709 
 710 test "semantic builder borrows a reusable compiler context" {
 711     const allocator = std.testing.allocator;
 712 
 713     var ctx = try buildSemanticContext(allocator, ir.Context.Limits.testing);
 714     defer ctx.deinit(allocator);
 715     const baseline = SemanticResourceCounts.capture(&ctx);
 716 
 717     var abandoned_builder = try Builder.initBorrowing(allocator, &ctx);
 718     errdefer abandoned_builder.deinit();
 719     try std.testing.expect(ctx.operationCount() > baseline.operations);
 720     abandoned_builder.deinit();
 721     try baseline.expectEqual(&ctx);
 722 
 723     var abandoned_iterate_builder = try Builder.initBorrowing(allocator, &ctx);
 724     errdefer abandoned_iterate_builder.deinit();
 725     const abandoned_type = try abandoned_iterate_builder.tensor(.i1, &.{4});
 726     var abandoned_function = try abandoned_iterate_builder.beginFunction(
 727         "borrowed_abandoned_iterate",
 728         &.{abandoned_type},
 729         &.{abandoned_type},
 730     );
 731     const abandoned_iterate = try abandoned_function.beginIterate(&.{abandoned_function.parameter(0)}, 4);
 732     _ = abandoned_iterate;
 733     abandoned_iterate_builder.deinit();
 734     try baseline.expectEqual(&ctx);
 735 
 736     var first_builder = try Builder.initBorrowing(allocator, &ctx);
 737     errdefer first_builder.deinit();
 738     const ty = try first_builder.tensor(.f32, &.{4});
 739     var first_function = try first_builder.beginFunction("borrowed_first", &.{ty}, &.{ty});
 740     try first_function.return_(&.{first_function.parameter(0)});
 741     try first_function.finish();
 742     const first_module = try first_builder.finish();
 743     var first_module_owned = true;
 744     defer if (first_module_owned) first_module.deinit();
 745     try std.testing.expect(!first_module.owns_ctx);
 746     first_module.deinit();
 747     first_module_owned = false;
 748     try baseline.expectEqual(&ctx);
 749 
 750     var second_builder = try Builder.initBorrowing(allocator, &ctx);
 751     errdefer second_builder.deinit();
 752     const second_ty = try second_builder.tensor(.i1, &.{4});
 753     var second_function = try second_builder.beginFunction("borrowed_second", &.{second_ty}, &.{second_ty});
 754     var second_iterate = try second_function.beginIterate(&.{second_function.parameter(0)}, 4);
 755     try second_iterate.yield_(second_iterate.carry(0), &.{second_iterate.carry(0)});
 756     try second_function.return_(&.{second_iterate.result(0)});
 757     try second_function.finish();
 758     const second_module = try second_builder.finish();
 759     var second_module_owned = true;
 760     defer if (second_module_owned) second_module.deinit();
 761 
 762     try second_module.verify();
 763     try std.testing.expect(ctx.isDialectLoaded("accy"));
 764     try std.testing.expect(ctx.isDialectLoaded("func"));
 765     second_module.deinit();
 766     second_module_owned = false;
 767     try baseline.expectEqual(&ctx);
 768 }
 769 
 770 test "semantic builder erases operations rejected before insertion" {
 771     const allocator = std.testing.allocator;
 772     var ctx = try buildSemanticContext(allocator, ir.Context.Limits.testing);
 773     defer ctx.deinit(allocator);
 774     const baseline = SemanticResourceCounts.capture(&ctx);
 775 
 776     var builder = try Builder.initBorrowing(allocator, &ctx);
 777     defer builder.deinit();
 778     const narrow = try builder.tensor(.f32, &.{4});
 779     const wide = try builder.tensor(.f32, &.{8});
 780     var function = try builder.beginFunction("borrowed_rejected_op", &.{ narrow, wide }, &.{narrow});
 781     const before_rejection = SemanticResourceCounts.capture(&ctx);
 782     if (function.add(function.parameter(0), function.parameter(1))) |_| {
 783         return error.TestExpectedError;
 784     } else |_| {}
 785     try before_rejection.expectEqual(&ctx);
 786 
 787     try function.return_(&.{function.parameter(0)});
 788     try function.finish();
 789     {
 790         const module = try builder.finish();
 791         defer module.deinit();
 792         try module.verify();
 793     }
 794     try baseline.expectEqual(&ctx);
 795 }
 796 
 797 fn checkBorrowedBuilderAllocationFailures(allocator: std.mem.Allocator) !void {
 798     var ctx = try buildSemanticContext(allocator, ir.Context.Limits.testing);
 799     defer ctx.deinit(allocator);
 800     const baseline = SemanticResourceCounts.capture(&ctx);
 801     defer baseline.expectEqual(&ctx) catch unreachable;
 802 
 803     var builder = try Builder.initBorrowing(allocator, &ctx);
 804     defer builder.deinit();
 805     const ty = try builder.tensor(.f32, &.{4});
 806     var function = try builder.beginFunction("borrowed_allocation_failure", &.{ ty, ty }, &.{ty});
 807     const sum = try function.add(function.parameter(0), function.parameter(1));
 808     try function.return_(&.{sum});
 809     try function.finish();
 810     const module = try builder.finish();
 811     defer module.deinit();
 812     try module.verify();
 813 }
 814 
 815 test "borrowed semantic builder restores context resources on allocation failure" {
 816     try @import("../fixture/root.zig").checkAllAllocationFailures(
 817         checkBorrowedBuilderAllocationFailures,
 818         .{},
 819     );
 820 }
 821 
 822 test "semantic builder preserves explicit source locations" {
 823     const allocator = std.testing.allocator;
 824     const module_location = ir.Location.getFile("model.chiclet", 1, 1);
 825     const function_location = ir.Location.getFile("model.chiclet", 3, 1);
 826     const add_location = ir.Location.getFile("model.chiclet", 4, 7);
 827     const return_location = ir.Location.getFile("model.chiclet", 5, 3);
 828 
 829     var builder = try Builder.initAt(allocator, Builder.ContextLimits.testing, module_location);
 830     errdefer builder.deinit();
 831     const ty = try builder.tensor(.f32, &.{8});
 832     var function = try builder.beginFunctionAt(
 833         "located_add",
 834         &.{ ty, ty },
 835         &.{ty},
 836         function_location,
 837     );
 838     function.setLocation(add_location);
 839     const sum = try function.add(function.parameter(0), function.parameter(1));
 840     const add_op: *ir.Operation = @ptrCast(@alignCast(sum.getDefiningOp().?));
 841     function.setLocation(return_location);
 842     try function.return_(&.{sum});
 843     const return_op: *ir.Operation = @ptrCast(@alignCast(function.entry.operations.tail.?));
 844     try function.finish();
 845     const module = try builder.finish();
 846     defer module.deinit();
 847 
 848     var module_operations = bodyBlock(module.choir_module).getOperations();
 849     const function_op = module_operations.next() orelse return error.TestExpectedResult;
 850     try std.testing.expect(module.choir_module.location.eql(module_location));
 851     try std.testing.expect(function_op.location.eql(function_location));
 852     try std.testing.expect(add_op.location.eql(add_location));
 853     try std.testing.expect(return_op.location.eql(return_location));
 854 }
 855 
 856 test "semantic builder exposes div operation" {
 857     const allocator = std.testing.allocator;
 858 
 859     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 860     defer builder.deinit();
 861 
 862     const ty = try builder.tensor(.f32, &.{4});
 863     var function = try builder.beginFunction("semantic_div", &.{ ty, ty }, &.{ty});
 864     const quotient = try function.div(function.parameter(0), function.parameter(1));
 865     try function.return_(&.{quotient});
 866     try function.finish();
 867 
 868     const module = try builder.finish();
 869     defer module.deinit();
 870 
 871     try module.verify();
 872 }
 873 
 874 test "semantic builder exposes exp operation" {
 875     const allocator = std.testing.allocator;
 876 
 877     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 878     defer builder.deinit();
 879 
 880     const ty = try builder.tensor(.f32, &.{4});
 881     var function = try builder.beginFunction("semantic_exp", &.{ty}, &.{ty});
 882     const out = try function.exp(function.parameter(0));
 883     try function.return_(&.{out});
 884     try function.finish();
 885 
 886     const module = try builder.finish();
 887     defer module.deinit();
 888 
 889     try module.verify();
 890 }
 891 
 892 test "semantic builder exposes wider elementwise math operations" {
 893     const allocator = std.testing.allocator;
 894 
 895     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 896     defer builder.deinit();
 897 
 898     const ty = try builder.tensor(.f32, &.{4});
 899     var function = try builder.beginFunction("semantic_wide_elementwise", &.{ ty, ty }, &.{ty});
 900     const lo = try function.min(function.parameter(0), function.parameter(1));
 901     const hi = try function.max(function.parameter(0), function.parameter(1));
 902     const delta = try function.sub(hi, lo);
 903     const magnitude = try function.abs(delta);
 904     const root = try function.sqrt(magnitude);
 905     const logarithm = try function.log(root);
 906     const sine = try function.sin(logarithm);
 907     const cosine = try function.cos(sine);
 908     const tangent = try function.tan(cosine);
 909     const powered = try function.pow(tangent, root);
 910     const quadrant = try function.atan2(powered, delta);
 911     try function.return_(&.{quadrant});
 912     try function.finish();
 913 
 914     const module = try builder.finish();
 915     defer module.deinit();
 916 
 917     try module.verify();
 918 }
 919 
 920 test "semantic builder exposes activation operation" {
 921     const allocator = std.testing.allocator;
 922 
 923     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 924     defer builder.deinit();
 925 
 926     const ty = try builder.tensor(.f32, &.{8});
 927     var function = try builder.beginFunction("semantic_activation", &.{ty}, &.{ty});
 928     const out = try function.activation(function.parameter(0), .silu);
 929     try function.return_(&.{out});
 930     try function.finish();
 931 
 932     const module = try builder.finish();
 933     defer module.deinit();
 934 
 935     try module.verify();
 936 }
 937 
 938 test "semantic builder exposes einsum operation" {
 939     const allocator = std.testing.allocator;
 940 
 941     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 942     defer builder.deinit();
 943 
 944     const lhs_ty = try builder.tensor(.f32, &.{ 4, 8 });
 945     const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
 946     const out_ty = try builder.tensor(.f32, &.{ 4, 16 });
 947     var function = try builder.beginFunction("semantic_einsum", &.{ lhs_ty, rhs_ty }, &.{out_ty});
 948     const out = try function.einsum(&.{ function.parameter(0), function.parameter(1) }, out_ty, "ik,kj->ij");
 949     try function.return_(&.{out});
 950     try function.finish();
 951 
 952     const module = try builder.finish();
 953     defer module.deinit();
 954 
 955     try module.verify();
 956 }
 957 
 958 test "semantic builder exposes kernel_call operation" {
 959     const allocator = std.testing.allocator;
 960 
 961     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
 962     defer builder.deinit();
 963 
 964     const ty = try builder.tensor(.f32, &.{4});
 965     var function = try builder.beginFunction("semantic_kernel_call", &.{ty}, &.{ty});
 966     const call = try function.kernelCall(
 967         &.{function.parameter(0)},
 968         &.{ty},
 969         .{
 970             .target = "scale_f32",
 971             .operand_effects = &.{.none},
 972             .result_aliases = &.{null},
 973         },
 974     );
 975     try function.return_(&.{call.getResult(0).?});
 976     try function.finish();
 977 
 978     const module = try builder.finish();
 979     defer module.deinit();
 980 
 981     try module.verify();
 982 }
 983 
 984 test "function builder emits compare convert and select operations" {
 985     var builder = try Builder.init(std.testing.allocator, Builder.ContextLimits.testing);
 986     errdefer builder.deinit();
 987     const vec_ty = try builder.tensor(.f32, &.{8});
 988     const pred_ty = try builder.tensor(.i1, &.{8});
 989     const int_ty = try builder.tensor(.i32, &.{8});
 990     var fb = try builder.beginFunction("semantic_compare_convert_select", &.{ vec_ty, vec_ty }, &.{int_ty});
 991     const pred = try fb.compare(fb.parameter(0), fb.parameter(1), pred_ty, .lt);
 992     const chosen = try fb.select(pred, fb.parameter(0), fb.parameter(1));
 993     const out = try fb.convert(chosen, int_ty, .i32);
 994     try fb.return_(&.{out});
 995     try fb.finish();
 996     const module = try builder.finish();
 997     defer module.deinit();
 998     try module.verify();
 999 }
1000 
1001 test "semantic builder rejects shape-mismatched elementwise misuse at call site" {
1002     var builder = try Builder.init(std.testing.allocator, Builder.ContextLimits.testing);
1003     defer builder.deinit();
1004     const narrow_ty = try builder.tensor(.f32, &.{4});
1005     const wide_ty = try builder.tensor(.f32, &.{8});
1006     var fb = try builder.beginFunction("semantic_misuse_shape", &.{ narrow_ty, wide_ty }, &.{narrow_ty});
1007     try std.testing.expectError(error.ShapeMismatch, fb.add(fb.parameter(0), fb.parameter(1)));
1008 }
1009 
1010 test "semantic builder rejects dtype-mismatched elementwise misuse at call site" {
1011     var builder = try Builder.init(std.testing.allocator, Builder.ContextLimits.testing);
1012     defer builder.deinit();
1013     const f32_ty = try builder.tensor(.f32, &.{4});
1014     const i32_ty = try builder.tensor(.i32, &.{4});
1015     var fb = try builder.beginFunction("semantic_misuse_dtype", &.{ f32_ty, i32_ty }, &.{f32_ty});
1016     try std.testing.expectError(error.DTypeMismatch, fb.add(fb.parameter(0), fb.parameter(1)));
1017 }
1018 
1019 test "semantic builder rejects non-boolean select predicates at call site" {
1020     var builder = try Builder.init(std.testing.allocator, Builder.ContextLimits.testing);
1021     defer builder.deinit();
1022     const vec_ty = try builder.tensor(.f32, &.{4});
1023     const pred_wrong_ty = try builder.tensor(.i32, &.{4});
1024     var fb = try builder.beginFunction("semantic_misuse_select", &.{ vec_ty, pred_wrong_ty }, &.{vec_ty});
1025     try std.testing.expectError(error.DTypeMismatch, fb.select(fb.parameter(1), fb.parameter(0), fb.parameter(0)));
1026 }
1027 
1028 test "semantic builder rejects mismatched function returns" {
1029     const allocator = std.testing.allocator;
1030 
1031     var builder = try Builder.init(allocator, Builder.ContextLimits.testing);
1032     defer builder.deinit();
1033 
1034     const input_ty = try builder.tensor(.f32, &.{4});
1035     const result_ty = try builder.tensor(.f32, &.{8});
1036     var function = try builder.beginFunction("bad_return", &.{input_ty}, &.{result_ty});
1037 
1038     try std.testing.expectError(error.ResultTypeMismatch, function.return_(&.{function.parameter(0)}));
1039 }