cwenzi/neuroflow-cpp

1#include <iostream>2#include "../include/neuroflow/tensor.hpp"3 4using namespace neuroflow;5 6int main() {7    std::cout << "GEMM transpose test..." << std::endl;8    9    // 简化场景10    Tensor A({1, 64});  // input {batch, d_model}11    Tensor B({64, 64}); // weight {d_model, d_model}12    Tensor C({1, 64});  // output {batch, d_model}13    14    float* a = A.as_fp32();15    float* b = B.as_fp32();16    for (size_t i = 0; i < A.numel(); ++i) a[i] = 0.1f * i;17    for (size_t i = 0; i < B.numel(); ++i) b[i] = 0.01f * i;18    19    std::cout << "A shape: [" << A.shape_[0] << ", " << A.shape_[1] << "]" << std::endl;20    std::cout << "B shape: [" << B.shape_[0] << ", " << B.shape_[1] << "]" << std::endl;21    std::cout << "C shape: [" << C.shape_[0] << ", " << C.shape_[1] << "]" << std::endl;22    23    // transB=true 时的参数24    bool transA = false;25    bool transB = true;26    27    size_t M = transA ? A.shape_[1] : A.shape_[0];  // 128    size_t K = transA ? A.shape_[0] : A.shape_[1];  // 6429    size_t N = transB ? B.shape_[0] : B.shape_[1];  // 64 (B.shape_[0])30    31    std::cout << "M=" << M << ", K=" << K << ", N=" << N << std::endl;32    std::cout << "transA=" << transA << ", transB=" << transB << std::endl;33    34    std::cout << "Calling gemm..." << std::endl;35    TensorOps::gemm(A, B, C, transA, transB);36    37    std::cout << "C values: ";38    float* c = C.as_fp32();39    for (size_t i = 0; i < 5; ++i) std::cout << c[i] << " ";40    std::cout << std::endl;41    42    std::cout << "Success!" << std::endl;43    return 0;44}45