move feature space to gpu

src/loss_function/LossFunctionLogPearsonRMSE.cpp

@@ -64,19 +64,19 @@ void LossFunctionLogPearsonRMSE::set_a(const std::vector<int>& inds, int taskind
     }
 }
 
-void LossFunctionLogPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature_indices,
+void LossFunctionLogPearsonRMSE::set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models,
                                        int taskind,
                                        int start)
 {
-    for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
+    for (int model_idx = 0; model_idx < models.extent(1); ++model_idx)
     {
-        for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
+        for (int feature_idx = 0; feature_idx < models.extent(0); ++feature_idx)
         {
-            double* val_ptr = node_value_arrs::get_d_matrix_ptr(feature_indices[model_idx][ff]) +
+            double* val_ptr = node_value_arrs::get_d_matrix_ptr(models(feature_idx, model_idx)) +
                               start;
             std::transform(val_ptr,
                            val_ptr + _task_sizes_train[taskind],
-                           &_a(0, ff, model_idx),
+                           &_a(0, feature_idx, model_idx),
                            [](double val) { return std::log(val); });
         }
     }

src/loss_function/LossFunctionLogPearsonRMSE.hpp

@@ -97,7 +97,7 @@ public:
      */
     void set_a(const std::vector<int>& inds, int taskind, int start);
 
-    void set_a(const std::vector<std::vector<int>>& feature_indices, int taskind, int start) override;
+    void set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models, int taskind, int start) override;
 
     /**
      * @brief Set the A matrix used for solving the least squares regression

src/loss_function/LossFunctionPearsonRMSE.cpp

@@ -86,6 +86,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::vector<double> prop_train,
       _c("c", _n_task),
       _training_properties("batched_scores", 0),
       _batched_scores("batched_scores"),
+      _models("models", _n_feat),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -114,6 +115,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::shared_ptr<LossFunction> o
       _c("c", _n_task),
       _training_properties("batched_scores", 0),
       _batched_scores("batched_scores"),
+      _models("models", _n_feat),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -142,6 +144,7 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(const LossFunctionPearsonRMSE&
       _c("c", _n_task),
       _training_properties("batched_scores", 0),
       _batched_scores("batched_scores"),
+      _models("models", _n_feat),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -218,6 +221,8 @@ void LossFunctionPearsonRMSE::set_nfeat(int n_feat)
     {
         _training_properties(material_idx) = _prop_train[material_idx];
     }
+
+    Kokkos::resize(_models, _n_feat);
 }
 
 void LossFunctionPearsonRMSE::reset_projection_prop(
@@ -324,17 +329,27 @@ Kokkos::View<double*> LossFunctionPearsonRMSE::operator()(
     const size_t batch_size = feature_indices.size();
     assert(batch_size <= MAX_BATCHES);
 
+    assert(feature_indices.size() <= _models.extent(1));
+    assert(feature_indices[0].size() == _models.extent(0));
+    for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
+    {
+        for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
+        {
+            _models(ff, model_idx) = feature_indices[model_idx][ff];
+        }
+    }
+
     int start = 0;
     for (int task_idx = 0; task_idx < _n_task; ++task_idx)
     {
         Kokkos::deep_copy(_a, 1.0);
-        set_a(feature_indices, task_idx, start);
+        set_a(_models, task_idx, start);
         set_b(task_idx, start);
         Kokkos::fence();
         least_squares(task_idx, start);
 
         Kokkos::deep_copy(_a, 1.0);
-        set_a(feature_indices, task_idx, start);
+        set_a(_models, task_idx, start);
         Kokkos::fence();
         set_prop_train_est(_estimated_training_properties, task_idx, start);
 
@@ -392,36 +407,25 @@ double LossFunctionPearsonRMSE::test_loss(const std::vector<model_node_ptr>& fea
 
 void LossFunctionPearsonRMSE::set_a(const std::vector<int>& inds, int taskind, int start)
 {
-    set_a(std::vector<std::vector<int>>(1, inds), taskind, start);
+    for (int ff = 0; ff < inds.size(); ++ff)
+    {
+        _models(ff, 0) = inds[ff];
+    }
+    set_a(_models, taskind, start);
 }
 
-void LossFunctionPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature_indices,
+void LossFunctionPearsonRMSE::set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models,
                                     int taskind,
                                     int start)
 {
-    Kokkos::View<int**, Kokkos::LayoutLeft> models(
-        "models", feature_indices[0].size(), feature_indices.size());
-    auto max_idx = 0;
-    for (int model_idx = 0; model_idx < feature_indices.size(); ++model_idx)
-    {
-        for (int ff = 0; ff < feature_indices[model_idx].size(); ++ff)
-        {
-            models(ff, model_idx) = feature_indices[model_idx][ff];
-            max_idx = std::max(max_idx, feature_indices[model_idx][ff]);
-        }
-    }
-
     assert(_descriptor_matrix.extent(0) >= _task_sizes_train[taskind] + start);
-    assert(_descriptor_matrix.extent(1) > max_idx);
     Kokkos::deep_copy(_a, 1.0);
     auto a = _a;
     auto descriptor_matrix = _descriptor_matrix;
 
     auto policy = Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
         {0, 0, 0},
-        {static_cast<size_t>(_task_sizes_train[taskind]),
-         feature_indices[0].size(),
-         feature_indices.size()});
+        {static_cast<size_t>(_task_sizes_train[taskind]), models.extent(0), models.extent(1)});
     auto kernel = KOKKOS_LAMBDA(const int sample_idx, const int feature_idx, const int model_idx)
     {
         a(sample_idx, feature_idx, model_idx) = descriptor_matrix(sample_idx + start,

src/loss_function/LossFunctionPearsonRMSE.hpp

@@ -56,6 +56,7 @@ protected:
     Kokkos::View<double*> _work;    //!< Work vector for dgels
     Kokkos::View<double*> _training_properties;
     Kokkos::View<double[MAX_BATCHES]> _batched_scores;
+    Kokkos::View<int*[MAX_BATCHES], Kokkos::LayoutLeft> _models;
 
     Kokkos::View<double**, Kokkos::LayoutLeft> _descriptor_matrix;
     Kokkos::View<double* [MAX_BATCHES], Kokkos::LayoutLeft> _estimated_training_properties;
@@ -175,7 +176,7 @@ public:
      */
     virtual void set_a(const std::vector<int>& inds, int taskind, int start);
 
-    virtual void set_a(const std::vector<std::vector<int>>& feature_indices, int taskind, int start);
+    virtual void set_a(Kokkos::View<int**, Kokkos::LayoutLeft> models, int taskind, int start);
 
     /**
      * @brief Set the A matrix used for solving the least squares regression