more kokkos views, less copies

src/loss_function/LossFunctionPearsonRMSE.cpp

@@ -84,6 +84,8 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::vector<double> prop_train,
       _a("A", _n_samp, _n_dim),
       _b("b", prop_train.size()),
       _c("c", _n_task),
+      _training_properties("batched_scores", 0),
+      _batched_scores("batched_scores"),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -110,6 +112,8 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(std::shared_ptr<LossFunction> o
       _a("A", _n_samp, _n_dim),
       _b("b", _prop_train.size()),
       _c("c", _n_task),
+      _training_properties("batched_scores", 0),
+      _batched_scores("batched_scores"),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -136,6 +140,8 @@ LossFunctionPearsonRMSE::LossFunctionPearsonRMSE(const LossFunctionPearsonRMSE&
       _a("A", _n_samp, _n_dim),
       _b("b", _prop_train.size()),
       _c("c", _n_task),
+      _training_properties("batched_scores", 0),
+      _batched_scores("batched_scores"),
       _lwork(0),
       _estimated_training_properties("estimated_training_properties", 0)
 {
@@ -205,6 +211,13 @@ void LossFunctionPearsonRMSE::set_nfeat(int n_feat)
     }
 
     Kokkos::resize(_estimated_training_properties, _prop_train_est.size());
+
+    Kokkos::resize(_training_properties, _prop_train.size());
+
+    for (int material_idx = 0; material_idx < _prop_train.size(); ++material_idx)
+    {
+        _training_properties(material_idx) = _prop_train[material_idx];
+    }
 }
 
 void LossFunctionPearsonRMSE::reset_projection_prop(
@@ -317,6 +330,7 @@ Kokkos::View<double*> LossFunctionPearsonRMSE::operator()(
         Kokkos::deep_copy(_a, 1.0);
         set_a(feature_indices, task_idx, start);
         set_b(task_idx, start);
+        Kokkos::fence();
         least_squares(task_idx, start);
 
         Kokkos::deep_copy(_a, 1.0);
@@ -327,17 +341,9 @@ Kokkos::View<double*> LossFunctionPearsonRMSE::operator()(
         start += _task_sizes_train[task_idx];
     }
 
-    Kokkos::View<double[MAX_BATCHES]> scores("score");
-    Kokkos::View<double*> training_properties("training_properties", _prop_train.size());
+    ::get_mean_squared_difference(_batched_scores, _training_properties, _estimated_training_properties);
 
-    for (int material_idx = 0; material_idx < _prop_train.size(); ++material_idx)
-    {
-        training_properties(material_idx) = _prop_train[material_idx];
-    }
-
-    ::get_mean_squared_difference(scores, training_properties, _estimated_training_properties);
-
-    return scores;
+    return _batched_scores;
 }
 
 double LossFunctionPearsonRMSE::operator()(const std::vector<model_node_ptr>& feats)
@@ -406,6 +412,7 @@ void LossFunctionPearsonRMSE::set_a(const std::vector<std::vector<int>>& feature
 
     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;
 

src/loss_function/LossFunctionPearsonRMSE.hpp

@@ -55,6 +55,7 @@ protected:
     Kokkos::View<double*> _scores;  //!< Vector used to temporarily store _projection_prop_sum
     Kokkos::View<double*> _work;    //!< Work vector for dgels
     Kokkos::View<double*> _training_properties;
+    Kokkos::View<double[MAX_BATCHES]> _batched_scores;
 
     Kokkos::View<double**, Kokkos::LayoutLeft> _descriptor_matrix;
     Kokkos::View<double* [MAX_BATCHES], Kokkos::LayoutLeft> _estimated_training_properties;