tiny.accy.kernel.logical

1pub const ActivationSelectionRequest = struct {2    dtype: DType,3    kind: ActivationKind,4    extent: u64,5};

1pub const SelectedActivationKernel = struct {2    kind: ActivationKind,3    descriptor: library.CatalogDescriptor,4};

1pub const EinsumKernelKind = enum {2    batched_matrix_product,3    dot_product,4    matrix_product,5    matrix_vector_product,6    outer_product,7    sum_reduction,8    transpose,9};

1pub const EinsumOperand = struct {2    indices: []const u8,3    dims: []const i64,4};

1pub const EinsumSchedule = union(enum) {2    batched_matrix_product: library.BatchedMatrixProductSchedule,3    matrix_product: library.MatrixProductSchedule,4    matrix_vector_product: library.MatrixVectorProductSchedule,5    outer_product: library.OuterProductSchedule,6};

1pub const EinsumSelectionRequest = struct {2    dtype: DType,3    inputs: []const EinsumOperand,4    output_indices: []const u8,5    output_dims: []const i64,6    schedule: ?EinsumSchedule = null,7};

1pub const OwnedSelectedEinsumKernel = struct {2    kind: EinsumKernelKind,3    descriptor: library.OwnedCatalogDescriptor,45    pub fn deinit(self: *OwnedSelectedEinsumKernel) void {6        self.descriptor.deinit();7        self.* = undefined;8    }9};

1pub const SelectedEinsumKernel = struct {2    kind: EinsumKernelKind,3    descriptor: library.CatalogDescriptor,4};

1pub const FusedKernelKind = union(enum) {2    vector: FusedVectorKind,3    matrix_product: FusedLinalgEpilogue,4    matrix_vector_product: FusedLinalgEpilogue,5    row_normalization: FusedRowNormalizationKind,6};

1pub const FusedMatrixProductSelectionRequest = struct {2    dtype: DType,3    epilogue: FusedLinalgEpilogue,4    lhs_indices: []const u8,5    rhs_indices: []const u8,6    bias_indices: []const u8,7    output_indices: []const u8,8    lhs_dims: []const i64,9    rhs_dims: []const i64,10    bias_dims: []const i64,11    output_dims: []const i64,12    schedule: ?FusedMatrixProductSchedule = null,13};

1pub const FusedMatrixVectorProductSelectionRequest = struct {2    dtype: DType,3    epilogue: FusedLinalgEpilogue,4    matrix_indices: []const u8,5    vector_indices: []const u8,6    bias_indices: []const u8,7    output_indices: []const u8,8    matrix_dims: []const i64,9    vector_dims: []const i64,10    bias_dims: []const i64,11    output_dims: []const i64,12    schedule: ?FusedMatrixVectorProductSchedule = null,13};

1pub const FusedRowNormalizationSelectionRequest = struct {2    dtype: DType,3    kind: FusedRowNormalizationKind,4    rows: u64,5    cols: u64,6    schedule: ?FusedRowNormalizationSchedule = null,7};

1pub const FusedSelectionRequest = union(enum) {2    vector: FusedVectorSelectionRequest,3    matrix_product: FusedMatrixProductSelectionRequest,4    matrix_vector_product: FusedMatrixVectorProductSelectionRequest,5    row_normalization: FusedRowNormalizationSelectionRequest,6};

1pub const FusedVectorSelectionRequest = struct {2    dtype: DType,3    kind: FusedVectorKind,4    extent: u64,5};

1pub const SelectedFusedKernel = struct {2    kind: FusedKernelKind,3    descriptor: library.CatalogDescriptor,4};

1pub const RowNormalizationSelectionRequest = struct {2    dtype: DType,3    kind: RowNormalizationKind,4    rows: u64,5    cols: u64,6    schedule: ?RowNormalizationSchedule = null,7};

1pub const SelectedRowNormalizationKernel = struct {2    kind: RowNormalizationKind,3    descriptor: library.CatalogDescriptor,4};

1pub const Axis = struct {2    name: []const u8,3    extent: u64,4};

1pub const Domain2D = struct {2    x: Axis,3    y: Axis,4};

1pub const Domain3D = struct {2    x: Axis,3    y: Axis,4    z: Axis,5};

1pub fn Family(comptime definition: anytype) type {2    return struct {3        const FamilySource: type = kernel.Family(FamilyDefinition(definition));45        pub const name = definition.name;6        pub const parameters = definition.parameters;7        pub const Layout: type = FamilySource.Layout;8        pub const Instance: type = definition.Instance;9        pub const Limits: type = FamilySource.Limits;1011        pub const arg = FamilySource.arg;12        pub const schema = FamilySource.schema;13        pub const build = FamilySource.build;14        pub const buildNamed = FamilySource.buildNamed;15        pub const interpret = FamilySource.interpret;16        pub const launch = FamilySource.launch;17        pub const scheduleSnapshot = FamilySource.scheduleSnapshot;18        pub const createPlan = FamilySource.createPlan;19        pub const createCheckedPlan = FamilySource.createCheckedPlan;20        pub const compileFragment = FamilySource.compileFragment;21        pub const createKernelArtifact = FamilySource.createKernelArtifact;22        pub const createKernelCallArtifact = FamilySource.createKernelCallArtifact;23        pub const runCpu = FamilySource.runCpu;24        pub const runCpuWithDiagnostic = FamilySource.runCpuWithDiagnostic;25        pub const verify = FamilySource.verify;26    };27}

1pub fn Program(comptime definition: anytype) type {2    return struct {3        const Default: type = Scheduled(definition, schedule_mod.default());45        pub const name = definition.name;6        pub const parameters = definition.parameters;7        pub const Layout: type = Default.Layout;8        pub const Limits: type = Default.Limits;910        pub fn withSchedule(comptime schedule: anytype) type {11            return Scheduled(definition, schedule);12        }1314        pub const arg = Default.arg;15        pub const schema = Default.schema;16        pub const build = Default.build;17        pub const interpret = Default.interpret;18        pub const analyze = Default.analyze;19        pub const transform = Default.transform;20        pub const launch = Default.launch;21        pub const scheduleSnapshot = Default.scheduleSnapshot;22        pub const createPlan = Default.createPlan;23        pub const createCheckedPlan = Default.createCheckedPlan;24        pub const compileFragment = Default.compileFragment;25        pub const createKernelArtifact = Default.createKernelArtifact;26        pub const createKernelCallArtifact = Default.createKernelCallArtifact;27        pub const runCpu = Default.runCpu;28        pub const runCpuWithDiagnostic = Default.runCpuWithDiagnostic;29        pub const verify = Default.verify;30    };31}

1pub fn selectCatalog(request: ActivationSelectionRequest) ?SelectedActivationKernel {2    const descriptor = library.select(.{ .activation = .{3        .dtype = request.dtype,4        .kind = request.kind,5        .extent = request.extent,6    } }) orelse return null;7    return .{8        .kind = request.kind,9        .descriptor = descriptor,10    };11}

1pub fn selectCatalog(request: EinsumSelectionRequest) ?SelectedEinsumKernel {2    if (selectReduction(request)) |selected| return selected;3    if (selectBatchedMatrixProduct(request)) |selected| return selected;4    if (selectMatrixProduct(request)) |selected| return selected;5    if (selectMatrixVectorProduct(request)) |selected| return selected;6    if (selectOuterProduct(request)) |selected| return selected;7    if (selectTranspose(request)) |selected| return selected;8    return null;9}

1pub fn selectOwnedCatalog(backing_allocator: std.mem.Allocator, request: EinsumSelectionRequest) !?OwnedSelectedEinsumKernel {2    if (selectCatalog(request)) |selected| return ownedSelectedStatic(selected);3    if (try selectOwnedBatchedMatrixProduct(backing_allocator, request)) |selected| return selected;4    if (try selectOwnedMatrixProduct(backing_allocator, request)) |selected| return selected;5    if (try selectOwnedMatrixVectorProduct(backing_allocator, request)) |selected| return selected;6    if (try selectOwnedOuterProduct(backing_allocator, request)) |selected| return selected;7    return null;8}

1pub fn selectCatalog(request: FusedSelectionRequest) ?SelectedFusedKernel {2    return switch (request) {3        .vector => |vector| selectVector(vector),4        .matrix_product => |matrix_product| selectMatrixProduct(matrix_product),5        .matrix_vector_product => |matrix_vector_product| selectMatrixVectorProduct(matrix_vector_product),6        .row_normalization => |row_normalization| selectRowNormalization(row_normalization),7    };8}

1pub fn selectCatalog(request: RowNormalizationSelectionRequest) ?SelectedRowNormalizationKernel {2    const descriptor = library.select(.{ .row_normalization = .{3        .dtype = request.dtype,4        .kind = request.kind,5        .rows = request.rows,6        .cols = request.cols,7        .schedule = request.schedule,8    } }) orelse return null;9    return .{10        .kind = request.kind,11        .descriptor = descriptor,12    };13}

1pub fn Builder(comptime Inner: type, comptime Schedule: type) type {2    return struct {3        token: *anyopaque,4        schedule: Schedule,56        const Self = @This();78        pub fn init(source: *Inner, schedule: Schedule) Self {9            return .{10                .token = source,11                .schedule = schedule,12            };13        }1415        pub fn index1D(self: *Self, name: []const u8, extent: u64) !kernel_domain.Index1D {16            return self.schedule.index1D(self.owner(), domain.axis(name, extent));17        }1819        pub fn index2D(self: *Self, shape: domain.Domain2D) !kernel_domain.Index2D {20            return self.schedule.index2D(self.owner(), shape);21        }2223        pub fn index3D(self: *Self, shape: domain.Domain3D) !kernel_domain.Index3D {24            return self.schedule.index3D(self.owner(), shape);25        }2627        pub fn guardDo(self: *Self, condition: anytype, context: anytype, comptime body: anytype) !void {28            var active_guard = try self.owner().guard(condition);29            errdefer active_guard.abort();30            try body(self, context);31            try active_guard.leave();32        }3334        pub fn guardIndexDo(self: *Self, index: kernel_domain.Index1D, context: anytype, comptime body: anytype) !void {35            var active_guard = try self.owner().guardIndex(index);36            errdefer active_guard.abort();37            try body(self, index, context);38            try active_guard.leave();39        }4041        pub fn guardIndex2DDo(self: *Self, index: kernel_domain.Index2D, context: anytype, comptime body: anytype) !void {42            var x_guard = try self.owner().guardIndex(index.x);43            errdefer x_guard.abort();44            var y_guard = try self.owner().guardIndex(index.y);45            errdefer y_guard.abort();46            try body(self, index, context);47            try y_guard.leave();48            try x_guard.leave();49        }5051        pub fn guardIndex3DDo(self: *Self, index: kernel_domain.Index3D, context: anytype, comptime body: anytype) !void {52            var x_guard = try self.owner().guardIndex(index.x);53            errdefer x_guard.abort();54            var y_guard = try self.owner().guardIndex(index.y);55            errdefer y_guard.abort();56            var z_guard = try self.owner().guardIndex(index.z);57            errdefer z_guard.abort();58            try body(self, index, context);59            try z_guard.leave();60            try y_guard.leave();61            try x_guard.leave();62        }6364        pub fn forEach1D(self: *Self, name: []const u8, extent: u64, context: anytype, comptime body: anytype) !kernel_domain.Index1D {65            const index = try self.index1D(name, extent);66            try self.guardIndexDo(index, context, body);67            return index;68        }6970        pub fn forEach2D(self: *Self, shape: domain.Domain2D, context: anytype, comptime body: anytype) !kernel_domain.Index2D {71            const index = try self.index2D(shape);72            try self.guardIndex2DDo(index, context, body);73            return index;74        }7576        pub fn forEach3D(self: *Self, shape: domain.Domain3D, context: anytype, comptime body: anytype) !kernel_domain.Index3D {77            const index = try self.index3D(shape);78            try self.guardIndex3DDo(index, context, body);79            return index;80        }8182        pub fn fold(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, initial: anytype, context: anytype, comptime body: anytype) !@TypeOf(initial) {83            return self.owner().fold(lower, upper, step, initial, .{84                .logical = self,85                .payload = context,86            }, struct {87                fn each(_: anytype, iteration: anytype, acc: anytype, fold_ctx: anytype) !@TypeOf(acc) {88                    return body(fold_ctx.logical, iteration, acc, fold_ctx.payload);89                }90            }.each);91        }9293        pub fn foldRange(self: *Self, lower: i64, upper: i64, step: i64, initial: anytype, context: anytype, comptime body: anytype) !@TypeOf(initial) {94            const lower_value = try self.constantIndex(lower);95            const upper_value = try self.constantIndex(upper);96            const step_value = try self.constantIndex(step);97            return self.fold(lower_value, upper_value, step_value, initial, context, body);98        }99100        pub fn whileLoop(self: *Self, initial: anytype, context: anytype, comptime condition: anytype, comptime body: anytype) !@TypeOf(initial) {101            return self.owner().whileLoop(initial, .{102                .logical = self,103                .payload = context,104            }, struct {105                fn keepGoing(_: anytype, carry: anytype, loop_ctx: anytype) !base.Value {106                    return typed.raw(try condition(loop_ctx.logical, carry, loop_ctx.payload));107                }108            }.keepGoing, struct {109                fn each(_: anytype, carry: anytype, loop_ctx: anytype) !@TypeOf(carry) {110                    return body(loop_ctx.logical, carry, loop_ctx.payload);111                }112            }.each);113        }114115        pub fn forScope(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, init_args: []const base.Value, result_types: []const base.Type) !ForScope(Inner) {116            return self.owner().forScope(lower, upper, step, init_args, result_types);117        }118119        pub fn whileScope(self: *Self, init_args: []const base.Value, result_types: []const base.Type) !WhileScope(Inner) {120            return self.owner().whileScope(init_args, result_types);121        }122123        pub fn forDo(self: *Self, lower: base.Value, upper: base.Value, step: base.Value, context: anytype, comptime body: anytype) !base.For {124            var scope = try self.forScope(lower, upper, step, &.{}, &.{});125            errdefer scope.abort();126            try body(self, scope.inductionVar(), context);127            try scope.leave(&.{});128            return scope.loop;129        }130131        pub fn forValueDo(self: *Self, comptime dtype: DType, lower: base.Value, upper: base.Value, step: base.Value, context: anytype, comptime body: anytype) !base.For {132            var scope = try self.forScope(lower, upper, step, &.{}, &.{});133            errdefer scope.abort();134            try body(self, try self.castValue(scope.inductionVar(), dtype), context);135            try scope.leave(&.{});136            return scope.loop;137        }138139        pub fn forRangeDo(self: *Self, lower: i64, upper: i64, step: i64, context: anytype, comptime body: anytype) !base.For {140            const lower_value = try self.constantIndex(lower);141            const upper_value = try self.constantIndex(upper);142            const step_value = try self.constantIndex(step);143            return self.forDo(lower_value, upper_value, step_value, context, body);144        }145146        pub fn forRangeValueDo(self: *Self, comptime dtype: DType, lower: i64, upper: i64, step: i64, context: anytype, comptime body: anytype) !base.For {147            const lower_value = try self.constantIndex(lower);148            const upper_value = try self.constantIndex(upper);149            const step_value = try self.constantIndex(step);150            return self.forValueDo(dtype, lower_value, upper_value, step_value, context, body);151        }152153        pub fn argument(self: *Self, index: usize) base.Value {154            return self.owner().argument(index);155        }156157        pub fn globalId(self: *Self, dim: base.Dimension) !base.Value {158            return self.owner().globalId(dim);159        }160161        pub fn globalIdValue(self: *Self, comptime dtype: DType, dim: base.Dimension) !typed.Value(dtype) {162            return self.owner().globalIdValue(dtype, dim);163        }164165        pub fn laneId(self: *Self) !base.Value {166            return self.owner().laneId();167        }168169        pub fn threadId(self: *Self, dim: base.Dimension) !base.Value {170            return self.owner().threadId(dim);171        }172173        pub fn blockId(self: *Self, dim: base.Dimension) !base.Value {174            return self.owner().blockId(dim);175        }176177        pub fn blockDim(self: *Self, dim: base.Dimension) !base.Value {178            return self.owner().blockDim(dim);179        }180181        pub fn gridDim(self: *Self, dim: base.Dimension) !base.Value {182            return self.owner().gridDim(dim);183        }184185        pub fn warpReduce(self: *Self, op_kind: base.WarpOpKind, value: anytype) !@TypeOf(value) {186            return self.owner().warpReduce(op_kind, value);187        }188189        pub fn warpScan(self: *Self, op_kind: base.WarpOpKind, mode: base.WarpScanMode, value: anytype) !@TypeOf(value) {190            return self.owner().warpScan(op_kind, mode, value);191        }192193        pub fn warpId(self: *Self) !base.Value {194            return self.owner().warpId();195        }196197        pub fn mmaSync(self: *Self, shape: base.MmaShape, a: [4]base.Value, b: [2]base.Value, c: [4]base.Value) ![4]base.Value {198            return self.owner().mmaSync(shape, a, b, c);199        }200201        pub fn shuffleSync(self: *Self, mode: base.ShuffleMode, value: base.Value, lane_or_delta: base.Value) !base.Value {202            return self.owner().shuffleSync(mode, value, lane_or_delta);203        }204205        pub fn fence(self: *Self, scope: base.Scope) !void {206            return self.owner().fence(scope);207        }208209        pub fn asyncCopyShared(self: *Self, dst: base.Value, dst_index: base.Value, src: base.Value, src_index: base.Value, bytes: u32) !void {210            return self.owner().asyncCopyShared(dst, dst_index, src, src_index, bytes);211        }212213        pub fn asyncCopyCommit(self: *Self) !void {214            return self.owner().asyncCopyCommit();215        }216217        pub fn asyncCopyWait(self: *Self, groups: u32) !void {218            return self.owner().asyncCopyWait(groups);219        }220221        pub fn sharedBuffer(self: *Self, dtype: DType, size: u64) !base.Value {222            return self.owner().sharedBuffer(dtype, size);223        }224225        pub fn dynamicSharedBuffer(self: *Self, dtype: DType, size: u64, byte_offset: u64) !base.Value {226            return self.owner().dynamicSharedBuffer(dtype, size, byte_offset);227        }228229        pub fn barrier(self: *Self, scope: base.Scope) !void {230            return self.owner().barrier(scope);231        }232233        pub fn typedArgument(self: *Self, comptime dtype: DType, index: usize) typed.Value(dtype) {234            return self.owner().typedArgument(dtype, index);235        }236237        pub fn typedValue(self: *Self, comptime dtype: DType, value: base.Value) typed.Value(dtype) {238            return self.owner().typedValue(dtype, value);239        }240241        pub fn castValue(self: *Self, input: anytype, comptime dtype: DType) !typed.Value(dtype) {242            return self.owner().castValue(input, dtype);243        }244245        pub fn bufferArgument(self: *Self, comptime dtype: DType, index: usize) view.BufferView(dtype) {246            return self.owner().bufferArgument(dtype, index);247        }248249        pub fn bufferView(self: *Self, comptime dtype: DType, value: base.Value) view.BufferView(dtype) {250            return self.owner().bufferView(dtype, value);251        }252253        pub fn constantInt(self: *Self, dtype: DType, value: i64) !base.Value {254            return self.owner().constantInt(dtype, value);255        }256257        pub fn constantIndex(self: *Self, value: i64) !base.Value {258            return self.owner().constantIndex(value);259        }260261        pub fn constantFloat(self: *Self, dtype: DType, value: f64) !base.Value {262            return self.owner().constantFloat(dtype, value);263        }264265        pub fn constantBool(self: *Self, value: bool) !base.Value {266            return self.owner().constantBool(value);267        }268269        pub fn constantValue(self: *Self, comptime dtype: DType, value: anytype) !typed.Value(dtype) {270            return self.owner().constantValue(dtype, value);271        }272273        pub fn add(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {274            return self.owner().add(lhs, rhs);275        }276277        pub fn sub(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {278            return self.owner().sub(lhs, rhs);279        }280281        pub fn mul(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {282            return self.owner().mul(lhs, rhs);283        }284285        pub fn umulhi(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {286            return self.owner().umulhi(lhs, rhs);287        }288289        pub fn div(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {290            return self.owner().div(lhs, rhs);291        }292293        pub fn min(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {294            return self.owner().min(lhs, rhs);295        }296297        pub fn max(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {298            return self.owner().max(lhs, rhs);299        }300301        pub fn and_(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {302            return self.owner().and_(lhs, rhs);303        }304305        pub fn or_(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {306            return self.owner().or_(lhs, rhs);307        }308309        pub fn xor(self: *Self, lhs: base.Value, rhs: base.Value) !base.Value {310            return self.owner().xor(lhs, rhs);311        }312313        pub fn not(self: *Self, input: base.Value) !base.Value {314            return self.owner().not(input);315        }316317        pub fn popcount(self: *Self, input: base.Value) !base.Value {318            return self.owner().popcount(input);319        }320321        pub fn ballotSync(self: *Self, predicate: base.Value) !base.Value {322            return self.owner().ballotSync(predicate);323        }324325        pub fn shl(self: *Self, value: base.Value, shift: base.Value) !base.Value {326            return self.owner().shl(value, shift);327        }328329        pub fn shr(self: *Self, value: base.Value, shift: base.Value) !base.Value {330            return self.owner().shr(value, shift);331        }332333        pub fn ushr(self: *Self, value: base.Value, shift: base.Value) !base.Value {334            return self.owner().ushr(value, shift);335        }336337        pub fn bitcast(self: *Self, input: base.Value, dtype: DType) !base.Value {338            return self.owner().bitcast(input, dtype);339        }340341        pub fn neg(self: *Self, input: base.Value) !base.Value {342            return self.owner().neg(input);343        }344345        pub fn abs(self: *Self, input: base.Value) !base.Value {346            return self.owner().abs(input);347        }348349        pub fn sqrt(self: *Self, input: base.Value) !base.Value {350            return self.owner().sqrt(input);351        }352353        pub fn exp(self: *Self, input: base.Value) !base.Value {354            return self.owner().exp(input);355        }356357        pub fn log(self: *Self, input: base.Value) !base.Value {358            return self.owner().log(input);359        }360361        pub fn tanh(self: *Self, input: base.Value) !base.Value {362            return self.owner().tanh(input);363        }364365        pub fn sin(self: *Self, input: base.Value) !base.Value {366            return self.owner().sin(input);367        }368369        pub fn cos(self: *Self, input: base.Value) !base.Value {370            return self.owner().cos(input);371        }372373        pub fn tan(self: *Self, input: base.Value) !base.Value {374            return self.owner().tan(input);375        }376377        pub fn floor(self: *Self, input: base.Value) !base.Value {378            return self.owner().floor(input);379        }380381        pub fn round(self: *Self, input: base.Value) !base.Value {382            return self.owner().round(input);383        }384385        pub fn trunc(self: *Self, input: base.Value) !base.Value {386            return self.owner().trunc(input);387        }388389        pub fn tf32Round(self: *Self, input: base.Value) !base.Value {390            return self.owner().tf32Round(input);391        }392393        pub fn pow(self: *Self, base_value: base.Value, exponent: base.Value) !base.Value {394            return self.owner().pow(base_value, exponent);395        }396397        pub fn atan2(self: *Self, y: base.Value, x: base.Value) !base.Value {398            return self.owner().atan2(y, x);399        }400401        pub fn fma(self: *Self, a: base.Value, b: base.Value, c: base.Value) !base.Value {402            return self.owner().fma(a, b, c);403        }404405        pub fn compare(self: *Self, predicate: base.Compare, lhs: base.Value, rhs: base.Value) !base.Value {406            return self.owner().compare(predicate, lhs, rhs);407        }408409        pub fn select(self: *Self, condition: base.Value, true_value: base.Value, false_value: base.Value) !base.Value {410            return self.owner().select(condition, true_value, false_value);411        }412413        pub fn cast(self: *Self, input: base.Value, dtype: DType) !base.Value {414            return self.owner().cast(input, dtype);415        }416417        pub fn castIndex(self: *Self, input: base.Value) !base.Value {418            return self.owner().castIndex(input);419        }420421        pub fn linearIndex(self: *Self, index: anytype) !base.Value {422            return self.owner().linearIndex(index);423        }424425        pub fn linearIndexValue(self: *Self, index: anytype, comptime dtype: DType) !typed.Value(dtype) {426            return self.owner().linearIndexValue(index, dtype);427        }428429        pub fn loadIndex(self: *Self, memref: base.Value, index: anytype) !base.Value {430            return self.owner().loadIndex(memref, index);431        }432433        pub fn loadVector(self: *Self, memref: base.Value, index: base.Value, width: u32) !base.Value {434            return self.owner().loadVector(memref, index, width);435        }436437        pub fn extractLane(self: *Self, vec: base.Value, lane: u32, dtype: DType) !base.Value {438            return self.owner().extractLane(vec, lane, dtype);439        }440441        pub fn insertLane(self: *Self, vec: base.Value, scalar: base.Value, lane: u32) !base.Value {442            return self.owner().insertLane(vec, scalar, lane);443        }444445        pub fn packVector(self: *Self, lanes: [4]base.Value) !base.Value {446            return self.owner().packVector(lanes);447        }448449        pub fn storeIndex(self: *Self, value: base.Value, memref: base.Value, index: anytype) !void {450            try self.owner().storeIndex(value, memref, index);451        }452453        pub fn atomicRmw(454            self: *Self,455            kind: base.AtomicRmwKind,456            value: base.Value,457            memref: base.Value,458            index: base.Value,459        ) !base.Value {460            return self.owner().atomicRmw(kind, value, memref, index);461        }462463        pub fn atomicRmwIndex(464            self: *Self,465            kind: base.AtomicRmwKind,466            value: base.Value,467            memref: base.Value,468            index: anytype,469        ) !base.Value {470            return self.atomicRmw(kind, value, memref, try self.linearIndex(index));471        }472473        pub fn atomicCas(474            self: *Self,475            expected: base.Value,476            desired: base.Value,477            memref: base.Value,478            index: base.Value,479        ) !base.Value {480            return self.owner().atomicCas(expected, desired, memref, index);481        }482483        pub fn atomicCasIndex(484            self: *Self,485            expected: base.Value,486            desired: base.Value,487            memref: base.Value,488            index: anytype,489        ) !base.Value {490            return self.atomicCas(expected, desired, memref, try self.linearIndex(index));491        }492493        pub fn vec2(self: *Self, x: base.Value, y: base.Value) vector.Vec2 {494            return self.owner().vec2(x, y);495        }496497        pub fn typedVec2(self: *Self, comptime dtype: DType, x: anytype, y: anytype) !typed.Vec2(dtype) {498            return self.owner().typedVec2(dtype, x, y);499        }500501        pub fn vec3(self: *Self, x: base.Value, y: base.Value, z: base.Value) vector.Vec3 {502            return self.owner().vec3(x, y, z);503        }504505        pub fn typedVec3(self: *Self, comptime dtype: DType, x: anytype, y: anytype, z: anytype) !typed.Vec3(dtype) {506            return self.owner().typedVec3(dtype, x, y, z);507        }508509        pub fn splat2(self: *Self, value: base.Value) vector.Vec2 {510            return self.owner().splat2(value);511        }512513        pub fn typedSplat2(self: *Self, comptime dtype: DType, value: anytype) !typed.Vec2(dtype) {514            return self.owner().typedSplat2(dtype, value);515        }516517        pub fn splat3(self: *Self, value: base.Value) vector.Vec3 {518            return self.owner().splat3(value);519        }520521        pub fn typedSplat3(self: *Self, comptime dtype: DType, value: anytype) !typed.Vec3(dtype) {522            return self.owner().typedSplat3(dtype, value);523        }524525        fn owner(self: *Self) *Inner {526            return @ptrCast(@alignCast(self.token));527        }528    };529}

1pub fn wrap(inner: anytype, schedule: anytype) Builder(PointerChild(@TypeOf(inner)), @TypeOf(schedule)) {2    return Builder(PointerChild(@TypeOf(inner)), @TypeOf(schedule)).init(inner, schedule);3}

1pub fn axis(name: []const u8, extent: u64) Axis {2    return .{3        .name = name,4        .extent = extent,5    };6}

1const model = @import("../model/logical/root.zig");2const family_mod = @import("family.zig");3const program = @import("program.zig");4const selection = @import("selection/root.zig");56pub const Axis = model.Axis;7pub const ActivationKind = selection.ActivationKind;8pub const ActivationSelectionRequest = selection.ActivationSelectionRequest;9pub const Builder = model.Builder;10pub const Domain2D = model.Domain2D;11pub const Domain3D = model.Domain3D;12pub const EinsumKernelKind = selection.EinsumKernelKind;13pub const EinsumOperand = selection.EinsumOperand;14pub const EinsumSchedule = selection.EinsumSchedule;15pub const EinsumSelectionRequest = selection.EinsumSelectionRequest;16pub const Family = family_mod.Family;17pub const FusedKernelKind = selection.FusedKernelKind;18pub const FusedLinalgEpilogue = selection.FusedLinalgEpilogue;19pub const FusedMatrixProductSchedule = selection.FusedMatrixProductSchedule;20pub const FusedMatrixProductSelectionRequest = selection.FusedMatrixProductSelectionRequest;21pub const FusedMatrixVectorProductSchedule = selection.FusedMatrixVectorProductSchedule;22pub const FusedMatrixVectorProductSelectionRequest = selection.FusedMatrixVectorProductSelectionRequest;23pub const FusedRowNormalizationKind = selection.FusedRowNormalizationKind;24pub const FusedRowNormalizationSchedule = selection.FusedRowNormalizationSchedule;25pub const FusedRowNormalizationSelectionRequest = selection.FusedRowNormalizationSelectionRequest;26pub const FusedSelectionRequest = selection.FusedSelectionRequest;27pub const FusedVectorKind = selection.FusedVectorKind;28pub const FusedVectorSelectionRequest = selection.FusedVectorSelectionRequest;29pub const Program = program.Program;30pub const RowNormalizationKind = selection.RowNormalizationKind;31pub const RowNormalizationParameterization = selection.RowNormalizationParameterization;32pub const RowNormalizationSchedule = selection.RowNormalizationSchedule;33pub const RowNormalizationSelectionRequest = selection.RowNormalizationSelectionRequest;34pub const SelectedActivationKernel = selection.SelectedActivationKernel;35pub const SelectedEinsumKernel = selection.SelectedEinsumKernel;36pub const OwnedSelectedEinsumKernel = selection.OwnedSelectedEinsumKernel;37pub const SelectedFusedKernel = selection.SelectedFusedKernel;38pub const SelectedRowNormalizationKernel = selection.SelectedRowNormalizationKernel;39pub const axis = model.axis;40pub const schedule = model.schedule;41pub const selectActivationCatalog = selection.selectActivationCatalog;42pub const selectEinsumCatalog = selection.selectEinsumCatalog;43pub const selectOwnedEinsumCatalog = selection.selectOwnedEinsumCatalog;44pub const selectFusedCatalog = selection.selectFusedCatalog;45pub const selectRowNormalizationCatalog = selection.selectRowNormalizationCatalog;46pub const wrap = model.wrap;

1pub const logical = logical_mod;