BLAS Integration¶

fn detectAvailableLibrary() BlasLibrary {
    const target_os = @import("builtin").target.os.tag;
    return switch (target_os) {
        .macos, .ios => .accelerate,
        .linux, .windows => blk: {
            if (detectOpenBlas()) break :blk .openblas;
            if (detectMkl()) break :blk .mkl;
            break :blk .generic;
        },
        else => .generic,
    };
}

pub const BlasInterface = struct {
    vtable: *const VTable,
    context: *anyopaque,

    pub const VTable = struct {
        // Level 1 BLAS
        dot:  *const fn (*anyopaque, u32, []const f32, []const f32) f32,
        axpy: *const fn (*anyopaque, u32, f32, []const f32, []f32) void,
        scal: *const fn (*anyopaque, u32, f32, []f32) void,
        nrm2: *const fn (*anyopaque, u32, []const f32) f32,

        // Level 2 BLAS
        gemv: *const fn (*anyopaque, MatrixLayout, BlasOperation,
                         u32, u32, f32, []const f32, u32,
                         []const f32, f32, []f32) void,

        // Level 3 BLAS
        gemm: *const fn (*anyopaque, MatrixLayout, BlasOperation, BlasOperation,
                         u32, u32, u32, f32, []const f32, u32,
                         []const f32, u32, f32, []f32, u32) void,

        deinit: *const fn (*anyopaque) void,
    };
};

fn genericDot(context: *anyopaque, n: u32, x: []const f32, y: []const f32) f32 {
    _ = context;
    var result: f32 = 0.0;

    const simd_width = 8;  // AVX2: 8 x f32
    const simd_end = (n / simd_width) * simd_width;

    var i: u32 = 0;
    if (comptime std.simd.suggestVectorLength(f32)) |vec_len| {
        if (vec_len >= simd_width) {
            const Vec = @Vector(simd_width, f32);
            var sum_vec: Vec = @splat(0.0);
            while (i < simd_end) {
                const x_vec: Vec = x[i..][0..simd_width].*;
                const y_vec: Vec = y[i..][0..simd_width].*;
                sum_vec += x_vec * y_vec;
                i += simd_width;
            }
            // Horizontal reduction
            for (0..simd_width) |j| result += sum_vec[j];
        }
    }
    // Scalar tail
    while (i < n) : (i += 1) result += x[i] * y[i];
    return result;
}

fn genericGemm(context: *anyopaque, layout: MatrixLayout,
               transa: BlasOperation, transb: BlasOperation,
               m: u32, n: u32, k: u32,
               alpha: f32, a: []const f32, lda: u32,
               b: []const f32, ldb: u32,
               beta: f32, c: []f32, ldc: u32) void {
    // 1. Scale C by beta
    for (0..m) |i| {
        for (0..n) |j| {
            const idx = if (layout == .column_major) j * ldc + i
                        else i * ldc + j;
            c[idx] *= beta;
        }
    }
    // 2. Accumulate alpha * A * B
    for (0..m) |i| {
        for (0..n) |j| {
            var sum: f32 = 0.0;
            for (0..k) |l| {
                const a_idx = indexFor(layout, transa, i, l, lda);
                const b_idx = indexFor(layout, transb, l, j, ldb);
                sum += a[a_idx] * b[b_idx];
            }
            const c_idx = if (layout == .column_major) j * ldc + i
                          else i * ldc + j;
            c[c_idx] += alpha * sum;
        }
    }
}

pub fn detect() BlasConfig {
    const detected_library = detectAvailableLibrary();
    const cpu_count = @as(u32, @intCast(std.Thread.getCpuCount() catch 4));

    return BlasConfig{
        .library = detected_library,
        .num_threads = @max(1, cpu_count - 1),
        .use_threading = cpu_count > 1,
        .memory_alignment = 64,
        .prefer_column_major = true,
    };
}

pub fn init(allocator: std.mem.Allocator, config: BlasConfig) !BlasManager {
    const blas = switch (config.library) {
        .generic => blk: {
            var g = try GenericBlas.init(allocator, config);
            break :blk g.interface();
        },
        .openblas => blk: {
            var o = try OpenBlas.init(allocator, config);
            break :blk o.interface();
        },
        .mkl, .accelerate, .atlas => blk: {
            std.log.warn("{} not yet implemented, falling back to generic",
                         .{config.library});
            var g = try GenericBlas.init(allocator, config);
            break :blk g.interface();
        },
    };
    return BlasManager{ .blas = blas, .config = config,
                         .allocator = allocator, .stats = BlasStats.init() };
}

pub const BlasStats = struct {
    operation_counts: std.EnumMap(BlasOpType, u64),
    total_time_ns:    std.EnumMap(BlasOpType, u64),
    total_flops:      std.EnumMap(BlasOpType, u64),

    pub fn getGflops(self: *const BlasStats, op: BlasOpType) f64 {
        const flops = self.total_flops.get(op) orelse 0;
        const time_ns = self.total_time_ns.get(op) orelse 1;
        return @as(f64, @floatFromInt(flops)) / @as(f64, @floatFromInt(time_ns));
    }
};