matrix.cpp

#include "matrix.h"
#include "getCacheSize.h"
#include "pool.h"

// =============================================================================
// SIMD操作函数实现
// =============================================================================

namespace simd_ops {
    // 加载操作
    inline simd_f32 load(const float* ptr) {
#if defined(SIMD_ARCH_ARM_NEON)
        return vld1q_f32(ptr);
#elif defined(SIMD_ARCH_X86_SSE)
        #ifdef __AVX__
            return _mm256_loadu_ps(ptr);
        #else
            return _mm_loadu_ps(ptr);
        #endif
#elif defined(SIMD_ARCH_APPLE_METAL)
        return simd_make_float4(ptr[0], ptr[1], ptr[2], ptr[3]);
#else
        simd_f32 result;
        for (int i = 0; i < 4; ++i) result.data[i] = ptr[i];
        return result;
#endif
    }

    // 设置标量到所有通道
    inline simd_f32 set1(float value) {
#if defined(SIMD_ARCH_ARM_NEON)
        return vdupq_n_f32(value);
#elif defined(SIMD_ARCH_X86_SSE)
        #ifdef __AVX__
            return _mm256_set1_ps(value);
        #else
            return _mm_set1_ps(value);
        #endif
#elif defined(SIMD_ARCH_APPLE_METAL)
        return simd::float4(value);
#else
        simd_f32 result;
        for (int i = 0; i < 4; ++i) result.data[i] = value;
        return result;
#endif
    }

    // 加法
    inline simd_f32 add(simd_f32 a, simd_f32 b) {
#if defined(SIMD_ARCH_ARM_NEON)
        return vaddq_f32(a, b);
#elif defined(SIMD_ARCH_X86_SSE)
        #ifdef __AVX__
            return _mm256_add_ps(a, b);
        #else
            return _mm_add_ps(a, b);
        #endif
#elif defined(SIMD_ARCH_APPLE_METAL)
        return a + b;
#else
        simd_f32 result;
        for (int i = 0; i < 4; ++i) result.data[i] = a.data[i] + b.data[i];
        return result;
#endif
    }

    // 乘法
    inline simd_f32 mul(simd_f32 a, simd_f32 b) {
#if defined(SIMD_ARCH_ARM_NEON)
        return vmulq_f32(a, b);
#elif defined(SIMD_ARCH_X86_SSE)
        #ifdef __AVX__
            return _mm256_mul_ps(a, b);
        #else
            return _mm_mul_ps(a, b);
        #endif
#elif defined(SIMD_ARCH_APPLE_METAL)
        return a * b;
#else
        simd_f32 result;
        for (int i = 0; i < 4; ++i) result.data[i] = a.data[i] * b.data[i];
        return result;
#endif
    }

    // 乘加：a * b + c
    inline simd_f32 fmadd(simd_f32 a, simd_f32 b, simd_f32 c) {
#if defined(SIMD_ARCH_ARM_NEON)
        return vmlaq_f32(c, a, b);
#elif defined(SIMD_ARCH_X86_SSE) && defined(__FMA__)
        return _mm_fmadd_ps(a, b, c);
#elif defined(SIMD_ARCH_X86_SSE)
        return add(mul(a, b), c);
#elif defined(SIMD_ARCH_APPLE_METAL)
        return a * b + c;
#else
        simd_f32 result;
        for (int i = 0; i < 4; ++i) result.data[i] = a.data[i] * b.data[i] + c.data[i];
        return result;
#endif
    }

    // 水平求和
    inline float horizontal_sum(simd_f32 v) {
#if defined(SIMD_ARCH_ARM_NEON)
    // on aarch64 vaddvq_f32 performs a vector-wide horizontal add
#if defined(__aarch64__)
    return vaddvq_f32(v);
#else
    float32x2_t sum = vadd_f32(vget_low_f32(v), vget_high_f32(v));
    return vget_lane_f32(vpadd_f32(sum, sum), 0);
#endif
#elif defined(SIMD_ARCH_X86_SSE)
    #ifdef __AVX__
        // AVX version with 8 floats
        __m128 sum128 = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
        sum128 = _mm_hadd_ps(sum128, sum128);
        sum128 = _mm_hadd_ps(sum128, sum128);
        return _mm_cvtss_f32(sum128);
    #else
        // SSE version with 4 floats
        __m128 sum = _mm_hadd_ps(v, v);
        sum = _mm_hadd_ps(sum, sum);
        return _mm_cvtss_f32(sum);
    #endif
#elif defined(SIMD_ARCH_APPLE_METAL)
        return simd::reduce_add(v);
#else
        float sum = 0.0f;
        for (int i = 0; i < 4; ++i) sum += v.data[i];
        return sum;
#endif
    }
}

// =============================================================================
// 基础矩阵操作函数实现
// =============================================================================

// aligned allocation helper
void* aligned_alloc_helper(size_t align, size_t size) {
#if defined(_ISOC11_SOURCE) || (defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L)
    // aligned_alloc requires size to be multiple of align
    size_t sz = ((size + align - 1) / align) * align;
    return aligned_alloc(align, sz);
#else
    void* ptr = nullptr;
    size_t sz = ((size + align - 1) / align) * align;
    if (posix_memalign(&ptr, align, sz) != 0) return nullptr;
    return ptr;
#endif
}

// 向量点积
float vec_dot(const float* x, const float* y, int n) {
    float sum = 0.0f;
    for (int i = 0; i < n; ++i) {
        sum += x[i] * y[i];
    }
    return sum;
}

// 基础矩阵乘法（三重循环）
float* base_mul(const float* A, const float* B, float* C, int r, int k, int c, int bs) {
    if (!C) {
        C = static_cast<float*>(aligned_alloc_helper(64, sizeof(float) * r * c));
    }
    memset(C, 0, sizeof(float) * r * c);
    
    for (int i = 0; i < r; ++i) {
        for (int j = 0; j < c; ++j) {
            float sum = 0.0f;
            for (int p = 0; p < k; ++p) {
                sum += A[i * k + p] * B[p * c + j];
            }
            C[i * c + j] = sum;
        }
    }
    return C;
}

// 矩阵转置乘法
float* matrix_mul_trans(const float* A, const float* B, float* res, int r, int k, int c, int bs) {
    const size_t align = 64;
    float *b = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * c * k));
    if (!res)
        res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!b || !res) {
        std::cerr << "Aligned allocation failed" << std::endl;
        free(b);
        free(res);
        return nullptr;
    }
    memset(res, 0, sizeof(float) * r * c);
    for (int i = 0; i < c; ++i) {
        for (int j = 0; j < k; ++j) {
            b[i * k + j] = B[j * c + i];
        }
    }
    for (int i = 0; i < r; ++i) {
        for (int j = 0; j < c; ++j) {
            float sum = 0.0f;
            for (int p = 0; p < k; ++p) {
                sum += A[i * k + p] * b[j * k + p];
            }
            res[i * c + j] = sum;
        }
    }
    free(b);
    return res;
}

// 分块矩阵乘法
float* matrix_mul_block(const float* A, const float* B, float* res, int r, int k, int c, int block_size) {
    const size_t align = 64;
    if (!res)
        res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!res) {
        std::cerr << "Aligned allocation failed" << std::endl;
        return nullptr;
    }
    memset(res, 0, sizeof(float) * r * c);
    for (int i = 0; i < r; i += block_size) {
        for (int j = 0; j < c; j += block_size) {
            for (int p = 0; p < k; p += block_size) {
                int i_end = std::min(i + block_size, r);
                int j_end = std::min(j + block_size, c);
                int p_end = std::min(p + block_size, k);
                for (int ii = i; ii < i_end; ++ii) {
                    for (int jj = j; jj < j_end; ++jj) {
                        float sum = 0.0f;
                        for (int pp = p; pp < p_end; ++pp) {
                            sum += A[ii * k + pp] * B[pp * c + jj];
                        }
                        res[ii * c + jj] += sum;
                    }
                }
            }
        }
    }
    return res;
}

// 分块转置矩阵乘法
float* matrix_mul_trans_block(const float* A, const float* B, float* res, int r, int k, int c, int block_size) {
    const size_t align = 64;
    float *b = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * c * k));
    if (!res)
        res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!b || !res) {
        std::cerr << "Aligned allocation failed" << std::endl;
        free(b);
        free(res);
        return nullptr;
    }
    memset(res, 0, sizeof(float) * r * c);
    for (int i = 0; i < c; ++i) {
        for (int j = 0; j < k; ++j) {
            b[i * k + j] = B[j * c + i];
        }
    }
    for (int i = 0; i < r; i += block_size) {
        for (int j = 0; j < c; j += block_size) {
            for (int p = 0; p < k; p += block_size) {
                int i_end = std::min(i + block_size, r);
                int j_end = std::min(j + block_size, c);
                int p_end = std::min(p + block_size, k);
                for (int ii = i; ii < i_end; ++ii) {
                    for (int jj = j; jj < j_end; ++jj) {
                        float sum = 0.0f;
                        for (int pp = p; pp < p_end; ++pp) {
                            sum += A[ii * k + pp] * b[jj * k + pp];
                        }
                        res[ii * c + jj] += sum;
                    }
                }
            }
        }
    }
    free(b);
    return res;
}

// SIMD优化的分块转置矩阵乘法
float* matrix_mul_trans_block_with_simd(const float* A, const float* B, float* res, int r, int k, int c, int block_size) {
    const size_t align = 64;
    float *b = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * c * k));
    if (!res)
        res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!b || !res) {
        std::cerr << "Aligned allocation failed" << std::endl;
        free(b);
        free(res);
        return nullptr;
    }
    memset(res, 0, sizeof(float) * r * c);
    for (int i = 0; i < c; ++i) {
        for(int j = 0; j < k; ++j) {
            b[i * k + j] = B[j * c + i];
        }
    }
    
    for (int i = 0; i < r; i += block_size) {
        for (int j = 0; j < c; j += block_size) {
            for (int p = 0; p < k; p += block_size) {
                int i_end = std::min(i + block_size, r);
                int j_end = std::min(j + block_size, c);
                int p_end = std::min(p + block_size, k);
                
                for (int ii = i; ii < i_end; ++ii) {
                    for (int jj = j; jj < j_end; ++jj) {
                        float sum = 0.0f;
                        int pp;
                        
                        // 使用统一的SIMD类型和操作
                        simd_f32 sum_vec0 = simd_ops::set1(0.0f), sum_vec1 = simd_ops::set1(0.0f),
                                  sum_vec2 = simd_ops::set1(0.0f), sum_vec3 = simd_ops::set1(0.0f);
                        for (pp = p; pp + 3 * simd_ops::SIMD_WIDTH < p_end; pp += 4 * simd_ops::SIMD_WIDTH) {
                            simd_f32 a0 = simd_ops::load(&A[ii * k + pp]);
                            simd_f32 b0 = simd_ops::load(&b[jj * k + pp]);
                            sum_vec0 = simd_ops::fmadd(a0, b0, sum_vec0);

                            simd_f32 a1 = simd_ops::load(&A[ii * k + pp + simd_ops::SIMD_WIDTH]);
                            simd_f32 b1 = simd_ops::load(&b[jj * k + pp + simd_ops::SIMD_WIDTH]);
                            sum_vec1 = simd_ops::fmadd(a1, b1, sum_vec1);

                            simd_f32 a2 = simd_ops::load(&A[ii * k + pp + 2 * simd_ops::SIMD_WIDTH]);
                            simd_f32 b2 = simd_ops::load(&b[jj * k + pp + 2 * simd_ops::SIMD_WIDTH]);
                            sum_vec2 = simd_ops::fmadd(a2, b2, sum_vec2);

                            simd_f32 a3 = simd_ops::load(&A[ii * k + pp + 3 * simd_ops::SIMD_WIDTH]);
                            simd_f32 b3 = simd_ops::load(&b[jj * k + pp + 3 * simd_ops::SIMD_WIDTH]);
                            sum_vec3 = simd_ops::fmadd(a3, b3, sum_vec3);
                        }
                        
                        // 处理剩余的元素
                        for (; pp < p_end; ++pp) {
                            sum += A[ii * k + pp] * b[jj * k + pp];
                        }
                        
                        // 将SIMD结果加到标量和中
                        sum += simd_ops::horizontal_sum(sum_vec0);
                        sum += simd_ops::horizontal_sum(sum_vec1);
                        sum += simd_ops::horizontal_sum(sum_vec2);
                        sum += simd_ops::horizontal_sum(sum_vec3);
                        res[ii * c + jj] += sum;
                    }
                }
            }
        }
    }
    
    free(b);
    return res;
}

// 异步SIMD优化的分块转置矩阵乘法（多线程）
float* async_matrix_mul_trans_block_with_simd(const float* A, const float* B, float* res, int r, int k, int c, int block_size, int num_threads) {
    const size_t align = 64;
    // Transpose B into b as other functions do
    float *b = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * c * k));
    if (!res)
        res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!b || !res) {
        std::cerr << "Aligned allocation failed" << std::endl;
        free(b);
        free(res);
        return nullptr;
    }
    memset(res, 0, sizeof(float) * r * c);
    for (int i = 0; i < c; ++i) {
        for (int j = 0; j < k; ++j) {
            b[i * k + j] = B[j * c + i];
        }
    }

    // Determine number of worker threads to use
    if (num_threads <= 0) {
        unsigned int hw = std::thread::hardware_concurrency();
        num_threads = hw == 0 ? 4 : static_cast<int>(hw);
    }

    // Split work by rows: each task handles a contiguous range of rows [row_start, row_end)
    int rows_per_task = (r + num_threads - 1) / num_threads;
    std::vector<std::future<void>> futures;
    futures.reserve(num_threads);

    for (int t = 0; t < num_threads; ++t) {
        int row_start = t * rows_per_task;
        if (row_start >= r) break;
        int row_end = std::min(r, row_start + rows_per_task);

        // capture by value the pointers and ranges to avoid race on locals
        auto fut = ThreadPool::get_instance(num_threads).enqueue_task([=, &A, &b, &res]() {
            // For the assigned row range, perform blocked, transposed, simd-accelerated multiply
            for (int i = row_start; i < row_end; i += block_size) {
                for (int j = 0; j < c; j += block_size) {
                    for (int p = 0; p < k; p += block_size) {
                        int i_end = std::min(i + block_size, row_end);
                        int j_end = std::min(j + block_size, c);
                        int p_end = std::min(p + block_size, k);

                        for (int ii = i; ii < i_end; ++ii) {
                            for (int jj = j; jj < j_end; ++jj) {
                                float sum = 0.0f;
                                int pp;

                                simd_f32 sum_vec0 = simd_ops::set1(0.0f), sum_vec1 = simd_ops::set1(0.0f),
                                          sum_vec2 = simd_ops::set1(0.0f), sum_vec3 = simd_ops::set1(0.0f);
                                for (pp = p; pp + 3 * simd_ops::SIMD_WIDTH < p_end; pp += 4 * simd_ops::SIMD_WIDTH) {
                                    simd_f32 a0 = simd_ops::load(&A[ii * k + pp]);
                                    simd_f32 b0 = simd_ops::load(&b[jj * k + pp]);
                                    sum_vec0 = simd_ops::fmadd(a0, b0, sum_vec0);

                                    simd_f32 a1 = simd_ops::load(&A[ii * k + pp + simd_ops::SIMD_WIDTH]);
                                    simd_f32 b1 = simd_ops::load(&b[jj * k + pp + simd_ops::SIMD_WIDTH]);
                                    sum_vec1 = simd_ops::fmadd(a1, b1, sum_vec1);

                                    simd_f32 a2 = simd_ops::load(&A[ii * k + pp + 2 * simd_ops::SIMD_WIDTH]);
                                    simd_f32 b2 = simd_ops::load(&b[jj * k + pp + 2 * simd_ops::SIMD_WIDTH]);
                                    sum_vec2 = simd_ops::fmadd(a2, b2, sum_vec2);

                                    simd_f32 a3 = simd_ops::load(&A[ii * k + pp + 3 * simd_ops::SIMD_WIDTH]);
                                    simd_f32 b3 = simd_ops::load(&b[jj * k + pp + 3 * simd_ops::SIMD_WIDTH]);
                                    sum_vec3 = simd_ops::fmadd(a3, b3, sum_vec3);
                                }

                                for (; pp < p_end; ++pp) {
                                    sum += A[ii * k + pp] * b[jj * k + pp];
                                }

                                sum += simd_ops::horizontal_sum(sum_vec0);
                                sum += simd_ops::horizontal_sum(sum_vec1);
                                sum += simd_ops::horizontal_sum(sum_vec2);
                                sum += simd_ops::horizontal_sum(sum_vec3);
                                // Each task writes only to its own rows, so this is safe without locks
                                res[ii * c + jj] += sum;
                            }
                        }
                    }
                }
            }
        });

        futures.emplace_back(std::move(fut));
    }

    // wait for all futures to complete
    for (auto &f : futures) {
        if (f.valid()) f.get();
    }

    // ensure pool tasks drained (not strictly necessary since futures completed)
    ThreadPool::get_instance().wait_for_completion();

    free(b);
    return res;
}

// =============================================================================
// 工具函数实现
// =============================================================================

// 随机数生成
float rand_float(float s) {
    return 4 * s * (1 - s);
}

// 生成随机矩阵
float* random_matrix(int r, int c, float seed) {
    const size_t align = 64;
    float *res = static_cast<float*>(aligned_alloc_helper(align, sizeof(float) * r * c));
    if (!res) return nullptr;
    for (int i = 0; i < r * c; ++i) {
        res[i] = rand_float(seed);
    }
    return res;
}

// 生成测试矩阵
void matrix_gen(float *a, float *b, int N, float seed) {
    float s = seed;
    for(int i = 0; i < N * N; i++) {
        s = rand_float(s);
        a[i] = s;
        s = rand_float(s);
        b[i] = s;
    }
}

// 计算矩阵迹
float Trace(const float* A, int r, int c) {
    float sum = 0.0f;
    int n = std::min(r, c);
    for (int i = 0; i < n; ++i) {
        sum += A[i * c + i];
    }
    return sum;
}

// 打印矩阵
void print_matrix(const float* A, int r, int c) {
    for (int i = 0; i < r; ++i) {
        for (int j = 0; j < c; ++j) {
            std::cout << A[i * c + j] << " ";
        }
        std::cout << std::endl;
    }
}

// 矩阵比较
bool comp(float *a, float *b, int N) {
    for (int i = 0; i < N * N; ++i) {
        if (std::abs(a[i] - b[i]) > 1e-3) {
            std::cout << "Mismatch at index " << i << ": " << a[i] << " != " << b[i] << std::endl;
            return false;
        }
    }
    return true;
}

// =============================================================================
// 测试和解析函数实现
// =============================================================================

// 显示详细用法信息
inline void print_help() {
    std::cout << "Matrix Multiplication Performance Tester\n\n";
    std::cout << "Usage:\n";
    std::cout << "  ./program [OPTIONS] [SIZES...]\n\n";
    std::cout << "Options:\n";
    std::cout << "  -h, --help          Show this help message\n";
    std::cout << "  -t, --test          Run in test mode, with clear guidance on usage to check out best fit block size\n";
    std::cout << "  -s, --seed          Random seed\n";
    std::cout << "  -m, --methods LIST   Comma-separated list of methods to test\n";
    std::cout << "                       Available methods: base, block, transpose, transpose_block, simd, async_simd, best\n";
    std::cout << "                       Default: all methods, random seed 0.1f\n\n";
    std::cout << "Examples:\n";
    std::cout << "  ./program 512 1024          # Test all methods with sizes 512 and 1024\n";
    std::cout << "  ./program -m base,simd 512 # Test only base and simd methods with size 512\n";
    std::cout << "  ./program --help           # Show this help\n\n";
    std::cout << "Available methods:\n";
    std::cout << "  base            : Basic triple-loop matrix multiplication\n";
    std::cout << "  block           : Blocked matrix multiplication\n";
    std::cout << "  transpose       : Matrix multiplication with transposed B\n";
    std::cout << "  transpose_block : Blocked matrix multiplication with transposed B\n";
    std::cout << "  simd            : SIMD-optimized matrix multiplication\n";
    std::cout << "  async_simd      : Asynchronous SIMD-optimized matrix multiplication\n";
    #ifdef SIMD_ARCH_APPLE_METAL
    std::cout << "  my_optimized    : Apple M2 Pro optimized method using custom NEON kernel\n";
    std::cout << "  apple_accelerate : Apple M2 Pro optimized method using Accelerate framework\n";
    #endif
    #if defined(SIMD_ARCH_ARM_NEON)
    std::cout << "  neon_omp        : OpenMP + NEON optimized matrix multiplication\n";
    #endif
    std::cout << "  best            : Best optimized method for current hardware\n";
}

// 解析逗号分隔的方法列表
std::set<std::string> parse_methods(const std::string& method_list) {
    std::set<std::string> methods;
    std::stringstream ss(method_list);
    std::string method;
    
    while (std::getline(ss, method, ',')) {
        // 去除前后空格
        method.erase(0, method.find_first_not_of(" \t"));
        method.erase(method.find_last_not_of(" \t") + 1);
        if (!method.empty()) {
            methods.insert(method);
        }
    }
    return methods;
}

// 验证方法有效性
bool is_valid_method(const std::string& method) {
    static const std::set<std::string> valid_methods = {
        "base", "block", "transpose", "transpose_block", 
        "simd", "async_simd", "best"
        #ifdef SIMD_ARCH_APPLE_METAL
        , "my_optimized", "apple_accelerate"
        #endif
        #if defined(SIMD_ARCH_ARM_NEON)
        , "neon_omp"
        #endif
    };
    return valid_methods.find(method) != valid_methods.end();
}

// =============================================================================
// 平台特定优化函数实现
// =============================================================================

#ifdef SIMD_ARCH_X86_SSE
// AVX微内核（6x16）- 简化实现
void avx_micro_kernel_6x16(const float* A, const float* B, float* C, 
                           int i, int j, int p, int i_end, int j_end, int p_end,
                           int k, int ldc) {
    // 基础实现，实际使用时需要根据具体AVX指令优化
    for (int ii = i; ii < i_end; ++ii) {
        for (int jj = j; jj < j_end; ++jj) {
            float sum = 0.0f;
            for (int pp = p; pp < p_end; ++pp) {
                sum += A[ii * k + pp] * B[pp * ldc + jj];
            }
            C[ii * ldc + jj] += sum;
        }
    }
}

// Xeon E5优化版本
float* optimized_matrix_mul_xeon_e5(const float* A, const float* B, 
                                  float* res, int r, int k, int c, int block_size) {
    // 使用基础的SIMD优化实现
    return matrix_mul_trans_block_with_simd(A, B, res, r, k, c, block_size);
}
#endif

#ifdef SIMD_ARCH_APPLE_METAL
// Apple M2优化版本（使用Accelerate框架）
float* optimized_matrix_mul_apple_m2(const float* A, const float* B, float* C, int r, int k, int c, int bs) {
    // 使用Accelerate框架的简化版本
    return matrix_mul_trans_block_with_simd(A, B, C, r, k, c, 128);
}

// 自定义NEON优化版本
float* my_optimized_matrix_mul(const float* A, const float* B, float* res, int r, int k, int c, int block_size) {
    // 使用NEON优化的简化版本
    return matrix_mul_trans_block_with_simd(A, B, res, r, k, c, block_size);
}
#endif

#if defined(SIMD_ARCH_ARM_NEON) || defined(SIMD_ARCH_APPLE_METAL)
// ARM NEON优化版本
float* optimized_matrix_mul_arm_neon(const float* A, const float* B, float* C, int r, int k, int c) {
    return matrix_mul_trans_block_with_simd(A, B, C, r, k, c, 128);
}

// NEON + OpenMP优化版本
float* optimized_matrix_mul_neon_omp(const float* A, const float* B, float* C, int r, int k, int c, int block_size) {
    return async_matrix_mul_trans_block_with_simd(A, B, C, r, k, c, block_size);
}

// NEON辅助函数 - 简化实现
inline float horizontal_sum_neon(float32x4_t v) {
#if defined(__aarch64__)
    return vaddvq_f32(v);
#else
    float32x2_t sum = vadd_f32(vget_low_f32(v), vget_high_f32(v));
    return vget_lane_f32(vpadd_f32(sum, sum), 0);
#endif
}

inline void load_a_block_broadcast(float32x4_t a_vec[], const float* A, int i, int k, int p, int p_end) {
    // 简化实现
    for (int idx = 0; idx < (p_end - p + 3) / 4; ++idx) {
        a_vec[idx] = vdupq_n_f32(A[i * k + p + idx * 4]);
    }
}

inline void load_b_block(float32x4_t b_vec[], const float* B_transposed, int j, int k, int p, int p_end) {
    // 简化实现
    for (int idx = 0; idx < (p_end - p + 3) / 4; ++idx) {
        b_vec[idx] = vld1q_f32(&B_transposed[j * k + p + idx * 4]);
    }
}

inline void load_c_block(float32x4_t c_regs[][2], const float* C, int i, int j, int ldc) {
    // 简化实现
    c_regs[0][0] = vld1q_f32(&C[i * ldc + j]);
    c_regs[0][1] = vld1q_f32(&C[i * ldc + j + 4]);
}

inline void store_c_block(const float32x4_t c_regs[][2], float* C, int i, int j, int ldc) {
    // 简化实现
    vst1q_f32(&C[i * ldc + j], c_regs[0][0]);
    vst1q_f32(&C[i * ldc + j + 4], c_regs[0][1]);
}

inline void neon_outer_product_update(float32x4_t c_regs[][2], float32x4_t a_vec, float32x4_t b_vec[]) {
    // 简化实现
    c_regs[0][0] = vmlaq_f32(c_regs[0][0], a_vec, b_vec[0]);
    c_regs[0][1] = vmlaq_f32(c_regs[0][1], a_vec, b_vec[1]);
}

void aggressive_neon_kernel(const float* A, const float* B_transposed, float* C, int r, int k, int c) {
    // 简化的NEON内核实现
    const int block_size = 128;
    for (int i = 0; i < r; i += block_size) {
        for (int j = 0; j < c; j += block_size) {
            for (int p = 0; p < k; p += block_size) {
                int i_end = std::min(i + block_size, r);
                int j_end = std::min(j + block_size, c);
                int p_end = std::min(p + block_size, k);
                
                for (int ii = i; ii < i_end; ++ii) {
                    for (int jj = j; jj < j_end; ++jj) {
                        float sum = 0.0f;
                        for (int pp = p; pp < p_end; ++pp) {
                            sum += A[ii * k + pp] * B_transposed[jj * k + pp];
                        }
                        C[ii * c + jj] += sum;
                    }
                }
            }
        }
    }
}

// NEON微内核函数 - 简化实现
static inline void neon_micro_kernel_4x8(const float* A_block, const float* B_block, float* C_block, 
                                         int k, int ldc, int prefetch_offset) {
    // 简化实现
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 8; ++j) {
            float sum = 0.0f;
            for (int p = 0; p < k; ++p) {
                sum += A_block[i * k + p] * B_block[j * k + p];
            }
            C_block[i * ldc + j] += sum;
        }
    }
}

static inline void neon_micro_kernel_4x4(const float* A_block, const float* B_block, float* C_block, int k, int ldc) {
    // 简化实现
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            float sum = 0.0f;
            for (int p = 0; p < k; ++p) {
                sum += A_block[i * k + p] * B_block[j * k + p];
            }
            C_block[i * ldc + j] += sum;
        }
    }
}
#endif

// =============================================================================
// 通用优化函数实现
// =============================================================================

// 缓存优化的矩阵转置
void cache_optimized_transpose(float* dst, const float* src, int rows, int cols) {
    const int block_size = 64;
    for (int i = 0; i < rows; i += block_size) {
        for (int j = 0; j < cols; j += block_size) {
            int i_end = std::min(i + block_size, rows);
            int j_end = std::min(j + block_size, cols);
            for (int ii = i; ii < i_end; ++ii) {
                for (int jj = j; jj < j_end; ++jj) {
                    dst[jj * rows + ii] = src[ii * cols + jj];
                }
            }
        }
    }
}

// 自适应最佳性能函数
float* best_matrix_mul(const float* A, const float* B, float* C, int r, int k, int c, int bs) {
#if defined(SIMD_ARCH_APPLE_METAL)
    return optimized_matrix_mul_apple_m2(A, B, C, r, k, c, bs);
#elif defined(SIMD_ARCH_ARM_NEON)
    return optimized_matrix_mul_neon_omp(A, B, C, r, k, c, bs);
#elif defined(SIMD_ARCH_X86_SSE)
    return async_matrix_mul_trans_block_with_simd(A, B, C, r, k, c, bs);
#else
    return matrix_mul_trans_block_with_simd(A, B, C, r, k, c, bs);
#endif
}

// =============================================================================
// 测试模块实现
// =============================================================================

void test_mod(int argc, char** argv) {
    // 默认参数
    int num_threads = std::thread::hardware_concurrency();
    if (num_threads == 0) num_threads = 4;
    int block_size = 64;
    std::vector<int> test_sizes = {512, 1024, 2048, 4096};
    float seed = 0.3f;
    int times = 3;
    std::vector<int> thread_counts = {1, 2, 4, 8, 16};  // 测试的线程数列表
    
    // 解析命令行参数
    for (int i = 0; i < argc; ++i) {
        if ((strcmp(argv[i], "--threads") == 0 || strcmp(argv[i], "-t") == 0) && i + 1 < argc) {
            num_threads = atoi(argv[i + 1]);
            if (num_threads <= 0) {
                std::cerr << "Invalid thread count. Using default: " << std::thread::hardware_concurrency() << std::endl;
                num_threads = std::thread::hardware_concurrency();
            }
            ++i;
        }
        else if ((strcmp(argv[i], "--block-size") == 0 || strcmp(argv[i], "-b") == 0) && i + 1 < argc) {
            block_size = atoi(argv[i + 1]);
            if (block_size <= 0) {
                std::cerr << "Invalid block size. Using default: 64" << std::endl;
                block_size = 64;
            }
            ++i;
        }
        else if ((strcmp(argv[i], "--sizes") == 0 || strcmp(argv[i], "-s") == 0) && i + 1 < argc) {
            test_sizes.clear();
            std::string sizes_str = argv[i + 1];
            std::stringstream ss(sizes_str);
            std::string size_str;
            while (std::getline(ss, size_str, ',')) {
                int size = atoi(size_str.c_str());
                if (size > 0) {
                    test_sizes.push_back(size);
                }
            }
            if (test_sizes.empty()) {
                std::cerr << "Invalid sizes. Using default: 512,1024,2048,4096" << std::endl;
                test_sizes = {512, 1024, 2048, 4096};
            }
            ++i;
        }
        else if ((strcmp(argv[i], "--seed") == 0 || strcmp(argv[i], "-e") == 0) && i + 1 < argc) {
            seed = atof(argv[i + 1]);
            ++i;
        }
        else if ((strcmp(argv[i], "--times") == 0 || strcmp(argv[i], "-n") == 0) && i + 1 < argc) {
            times = atoi(argv[i + 1]);
            if (times <= 0) {
                std::cerr << "Invalid times value. Using default: 3" << std::endl;
                times = 3;
            }
            ++i;
        }
        else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            std::cout << "Usage: " << argv[0] << " [options]" << std::endl;
            std::cout << "Options:" << std::endl;
            std::cout << "  -t, --threads <num>     Set number of threads (default: hardware concurrency)" << std::endl;
            std::cout << "  -b, --block-size <size> Set block size (default: 64)" << std::endl;
            std::cout << "  -s, --sizes <list>      Set matrix sizes (comma-separated, default: 512,1024,2048,4096)" << std::endl;
            std::cout << "  -e, --seed <value>      Set random seed (default: 0.3)" << std::endl;
            std::cout << "  -n, --times <num>       Set number of runs per test (default: 3)" << std::endl;
            std::cout << "  -h, --help              Show this help message" << std::endl;
            return;
        }
    }
    
    // 输出测试配置
    std::cout << "=== async_simd Matrix Multiplication Performance Test ===" << std::endl;
    std::cout << "Method: async_matrix_mul_trans_block_with_simd" << std::endl;
    std::cout << "Block size: " << block_size << std::endl;
    std::cout << "Random seed: " << seed << std::endl;
    std::cout << "Runs per test: " << times << std::endl;
    std::cout << "Thread counts to test: ";
    for (int tc : thread_counts) std::cout << tc << " ";
    std::cout << std::endl;
    std::cout << "Matrix sizes: ";
    for (int size : test_sizes) std::cout << size << " ";
    std::cout << std::endl << std::endl;
    
    // 输出表头
    std::cout << "Matrix Size\tThreads\tTime (ms)\tGFlops\t\tTrace" << std::endl;
    std::cout << "-----------\t-------\t---------\t-------\t\t-----" << std::endl;
    
    // 对每个矩阵大小和线程数进行测试
    for (int size : test_sizes) {
        // 分配矩阵内存
        float *A = static_cast<float*>(aligned_alloc_helper(64, sizeof(float) * size * size));
        float *B = static_cast<float*>(aligned_alloc_helper(64, sizeof(float) * size * size));
        float *C = static_cast<float*>(aligned_alloc_helper(64, sizeof(float) * size * size));
        
        if (!A || !B || !C) {
            std::cerr << "Error: Failed to allocate memory for matrices of size " << size << std::endl;
            continue;
        }
        
        // 生成测试矩阵
        matrix_gen(A, B, size, seed);
        
        for (int threads : thread_counts) {
            // 只对4096矩阵测试所有线程数，其他大小只测试默认线程数
            if (size != 4096 && threads != num_threads) {
                continue;
            }
            
            double total_time = 0.0;
            float trace = 0.0f;
            
            // 运行多次取平均值
            for (int run = 0; run < times; ++run) {
                // 清零结果矩阵
                memset(C, 0, sizeof(float) * size * size);
                
                auto start = std::chrono::high_resolution_clock::now();
                
                // 调用async_simd方法，传递指定的线程数
                float* result = async_matrix_mul_trans_block_with_simd(A, B, C, size, size, size, block_size, threads);
                
                auto end = std::chrono::high_resolution_clock::now();
                double elapsed_ms = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.0;
                total_time += elapsed_ms;
                
                // 计算trace用于验证正确性
                if (run == 0) {  // 只在第一次运行时计算trace
                    trace = 0.0f;
                    for (int i = 0; i < size; ++i) {
                        trace += result[i * size + i];
                    }
                }
            }
            
            double avg_time_ms = total_time / times;
            double gflops = (2.0 * size * size * size) / (avg_time_ms * 1e6);
            
            // 输出结果，格式化为表格形式
            printf("%-11d\t%-7d\t%-9.2f\t%-7.2f\t\t%.2e\n", 
                   size, threads, avg_time_ms, gflops, trace);
        }
        
        // 释放内存
        free(A);
        free(B);
        free(C);
    }
    
    std::cout << std::endl << "Test completed." << std::endl;
}