Function PYBIND11_MODULE

ngscuda/python_ngscuda.cpp:19–293

}

PYBIND11_MODULE(_ngscuda, m) {

  InitCUDA(1);
  InitCuLinalg();
  InitApplyIntegrationPoints();

  m.def("InitCuLinalg", &InitCuLinalg, "Initializing cublas and cusparse.");
  
  py::class_<UnifiedVector, BaseVector, shared_ptr<UnifiedVector>>
    (m, "UnifiedVector", "UnifiedVector for CUDA applications", py::multiple_inheritance())
    
    .def(py::init([] (int size)
                  { 
                    return make_shared<UnifiedVector>(size); 
                  }))
    .def(py::init([] (const BaseVector &vec) 
                  {
                    return make_shared<UnifiedVector>(vec);
                  }))
    .def(py::init([] (py::array_t<double> bvec)
                  {
                    auto vec = bvec.template unchecked<1>();
                    shared_ptr<UnifiedVector> uv = make_shared<UnifiedVector>(vec.size());
                    FlatVector<double> fv = uv->FVDouble();
                    for (size_t i = 0; i < vec.size(); i++)
                      {
                        fv(i) = vec(i);
                      }
                    return uv;
                  }))

    .def("UpdateHost", &UnifiedVector::UpdateHost)
    .def("UpdateDevice", &UnifiedVector::UpdateDevice)
    .def_property_readonly("__cuda_array_interface__", [](UnifiedVector& self)
    {
        self.UpdateDevice();
        auto ptr = reinterpret_cast<uintptr_t>(self.DevData());
        py::dict cai;
        cai["version"] = 2;
        cai["shape"]   = py::make_tuple(self.Size());
        // "<f8" = little-endian float64
        cai["typestr"] = "<f8";
        // data: (ptr, readonly_flag)
        cai["data"] = py::make_tuple(ptr, false);
        // contiguous 1D, so no strides
        cai["strides"] = py::none();
        return cai;
    })
    .def_property_readonly("dev_ptr", [](UnifiedVector& self)
    {
        return reinterpret_cast<uintptr_t>(self.DevData());
    })
    ;


  py::class_<DevMatrix, BaseMatrix, shared_ptr<DevMatrix>>
    (m, "DevBaseMatrix", "device matrix for CUDA applications");