tiny.accy.preparation.shape

1pub const ShapeLayoutAnalysis = struct {2    allocator: std.mem.Allocator,3    tensors: std.AutoHashMap(*ir.Value, TensorInfo),4    tensor_value_count: usize = 0,5    accy_op_count: usize = 0,67    pub fn init(allocator: std.mem.Allocator) ShapeLayoutAnalysis {8        return .{9            .allocator = allocator,10            .tensors = std.AutoHashMap(*ir.Value, TensorInfo).init(allocator),11        };12    }1314    pub fn deinit(self: *ShapeLayoutAnalysis) void {15        var iter = self.tensors.valueIterator();16        while (iter.next()) |info| {17            info.deinit(self.allocator);18        }19        self.tensors.deinit();20        self.* = undefined;21    }2223    pub fn get(self: *const ShapeLayoutAnalysis, value: *ir.Value) ?TensorInfo {24        return self.tensors.get(value);25    }2627    fn recordValue(self: *ShapeLayoutAnalysis, value: *ir.Value) !void {28        const info = try tensorInfoFromType(self.allocator, value.type) orelse return;29        errdefer {30            var mutable = info;31            mutable.deinit(self.allocator);32        }33        self.tensors.putAssumeCapacityNoClobber(value, info);34        self.tensor_value_count += 1;35    }36};

1pub const TensorInfo = struct {2    dtype: choir_abi.DType,3    dims: []i64,4    element_count: ?u64,5    row_major_strides: ?[]u64,67    pub fn rank(self: TensorInfo) usize {8        return self.dims.len;9    }1011    pub fn hasStaticLayout(self: TensorInfo) bool {12        return self.element_count != null and self.row_major_strides != null;13    }1415    fn deinit(self: *TensorInfo, allocator: std.mem.Allocator) void {16        allocator.free(self.dims);17        if (self.row_major_strides) |strides| allocator.free(strides);18        self.* = undefined;19    }20};

1pub fn getShapeLayoutAnalysis(2    pass_ctx: *passes.PassContext,3    op: *ir.Operation,4) !*ShapeLayoutAnalysis {5    return ShapeLayoutAnalysisRegistration.get(pass_ctx, op);6}

1pub fn shapeLayoutPropagationPass() passes.Pass {2    return .{3        .name = shape_layout_pass_name,4        .description = shape_layout_pass_description,5        .run_fn = runShapeLayoutPropagationPass,6        .work_contract = .{7            .identity = .{ .name = shape_layout_pass_name, .version = 1 },8            .estimate = propagationWork,9        },10    };11}

1pub const shape_layout_analysis_name = "accy-choir-shape-layout";

1pub const shape_layout_pass_description =2    "Decode Accy Choir tensor shape and row-major layout facts";

1pub const shape_layout_analysis_descriptor = ShapeLayoutAnalysisRegistration.descriptor;

1pub const shape = @import("shape/root.zig");

1const pass = @import("pass.zig");23pub const shape_layout_analysis_name = pass.shape_layout_analysis_name;4pub const shape_layout_pass_name = pass.shape_layout_pass_name;5pub const shape_layout_pass_description = pass.shape_layout_pass_description;6pub const TensorInfo = pass.TensorInfo;7pub const ShapeLayoutAnalysis = pass.ShapeLayoutAnalysis;8pub const shape_layout_analysis_descriptor = pass.shape_layout_analysis_descriptor;9pub const getShapeLayoutAnalysis = pass.getShapeLayoutAnalysis;10pub const shapeLayoutPropagationPass = pass.shapeLayoutPropagationPass;