cwenzi/neuroflow-cpp
1
1#include <iostream>2#include "../include/neuroflow/tensor.hpp"3 4using namespace neuroflow;5 6int main() {7 std::cout << "GEMM detail test..." << std::endl;8 9 Tensor A({2, 64}); // input10 Tensor B({32, 64}); // weight11 Tensor C({2, 32}); // output12 13 std::cout << "A shape: [" << A.shape_[0] << ", " << A.shape_[1] << "]" << std::endl;14 std::cout << "B shape: [" << B.shape_[0] << ", " << B.shape_[1] << "]" << std::endl;15 std::cout << "C shape: [" << C.shape_[0] << ", " << C.shape_[1] << "]" << std::endl;16 17 // Fill data18 float* a = A.as_fp32();19 float* b = B.as_fp32();20 for (size_t i = 0; i < A.numel(); ++i) a[i] = 0.1f * i;21 for (size_t i = 0; i < B.numel(); ++i) b[i] = 0.01f * i;22 23 std::cout << "Calling gemm..." << std::endl;24 TensorOps::gemm(A, B, C); // C = A @ B^T ?25 26 std::cout << "C numel: " << C.numel() << std::endl;27 std::cout << "C data_size: " << C.data_size_ << std::endl;28 29 float* c = C.as_fp32();30 std::cout << "First 5 C values: ";31 for (size_t i = 0; i < 5; ++i) std::cout << c[i] << " ";32 std::cout << std::endl;33 34 std::cout << "After gemm, trying to allocate new tensor..." << std::endl;35 Tensor new_t({10});36 37 std::cout << "Success!" << std::endl;38 return 0;39}40 