python_ngsmetal.cpp

#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>


#include <comp.hpp>
#include <python_comp.hpp>

#include "ngsmetal.hpp"
#include "metal_vector.hpp"
#include "metal_sparsematrix.hpp"
#include "metal_btdtb.hpp"

using namespace ngsmetal;

namespace ngsmetal {
  extern void ComputeBenchmarks();
  extern void ComputeBenchmarksLoadCompute();
  extern void ComputeBenchmarksLoad();  
}

PYBIND11_MODULE(ngsmetal, m)
{
  cout << "Loading ngs-metal library" << endl;

  ngsmetal::InitNgsMetal();
  ngsmetal::InitVectorKernels();
  ngsmetal::InitSparseMatrixKernels();
  ngsmetal::InitMatMat();

  py::class_<MetalVector, BaseVector, shared_ptr<MetalVector>> (m, "MetalVector")
    .def(py::init<size_t>(), py::arg("size"))
    ;

  py::class_<MetalSparseMatrix, BaseMatrix, shared_ptr<MetalSparseMatrix>> (m, "MetalSparseMatrix")
    .def(py::init<const BaseSparseMatrix&>(), py::arg("mat"))
    ;

  py::class_<MetalBTDTBMatrix, BaseMatrix, shared_ptr<MetalBTDTBMatrix>> (m, "MetalBTDTBMatrix")
    .def(py::init<const BaseMatrix&>(), py::arg("mat"))
    ;



  
  m.def("TestMatMat", [](size_t n, size_t m, size_t k, int runs) {
    
    Matrix<float> a(n,k), b(k,m), c(n,m);

    /*
    if (n%32 != 0) throw Exception("n must be multiple of 32");
    if (m%32 != 0) throw Exception("m must be multiple of 32");
    if (k%32 != 0) throw Exception("k must be multiple of 32");
    
    for (int i = 0; i < n; i++)
      for (int j = 0; j < k; j++)
        a(i,j) = sin(i+j);
    for (int i = 0; i < k; i++)
      for (int j = 0; j < m; j++)
        b(i,j) = sin(i*j);
    */
    a = 1; b = 2;

    for (int i = 0; i < runs; i++)
      Mult (a,b,c);

    if (c.Height()*c.Width() < 10000)
      cout << "err = " << L2Norm(c-a*b) << endl;
    // cout << "c = " << c << endl;
  });

  m.def("ComputeBenchmarks", &ComputeBenchmarks);
  m.def("ComputeBenchmarksLoadCompute", &ComputeBenchmarksLoadCompute);
  m.def("ComputeBenchmarksLoad", &ComputeBenchmarksLoad);
}