Merge pull request #2 from dannys4/small_sigmoid_LA_gpu_fix

Small sigmoid la gpu fix

CMakeLists.txt

@@ -88,7 +88,7 @@ if(NOT Kokkos_FOUND)
         FetchContent_Declare(
             kokkos
             GIT_REPOSITORY https://github.com/kokkos/kokkos
-            GIT_TAG 3.7.00
+            GIT_TAG 4.2.00
             GIT_SHALLOW TRUE
         )
         FetchContent_MakeAvailable(kokkos)

MParT/ComposedMap.h

@@ -50,12 +50,12 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
                   start at index \f$\sum_{j=1}^{k-1} C_j\f$.
     */
     using ConditionalMapBase<MemorySpace>::SetCoeffs;
-    void SetCoeffs(Kokkos::View<const double*, Kokkos::HostSpace> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs) override;
+    using ConditionalMapBase<MemorySpace>::WrapCoeffs;
+    void SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs) override;
+    void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs) override;
 
     #if defined(MPART_ENABLE_GPU)
-    void SetCoeffs(Kokkos::View<const double*, Kokkos::DefaultExecutionSpace::memory_space> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs) override;
+    void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<MemorySpace, Kokkos::HostSpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs) override;
     #endif
 
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return maps_.at(i);}

MParT/MonotoneComponent.h

@@ -1073,9 +1073,9 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
         ExpansionType expansion;
         double nugget;
 
-        SingleEvaluator(double* workspace_, double* cache_, PointType pt_, CoeffType coeffs_, QuadratureType quad_, ExpansionType expansion_, double nugget_):
+        KOKKOS_FUNCTION SingleEvaluator(double* workspace_, double* cache_, PointType pt_, CoeffType coeffs_, QuadratureType quad_, ExpansionType expansion_, double nugget_):
             workspace(workspace_), cache(cache_), pt(pt_), coeffs(coeffs_), quad(quad_), expansion(expansion_), nugget(nugget_) {};
-        double operator()(double x) {
+        KOKKOS_INLINE_FUNCTION double operator()(double x) {
             return EvaluateSingle(cache, workspace, pt, x, coeffs, quad, expansion, nugget);
         }
     };

MParT/MultivariateExpansionWorker.h

@@ -336,7 +336,6 @@ class MultivariateExpansionWorker
     {
         const unsigned int numTerms = multiSet_.Size();
         double df = 0.0;
-        unsigned int posInd;
 
         for(int wrt=-1; wrt<int(dim_); ++wrt){
             if(wrt>=0){
@@ -529,7 +528,6 @@ class MultivariateExpansionWorker
     KOKKOS_FUNCTION double GetTermValMixedCoeffDeriv(unsigned int termInd, int wrt, const double* cache, unsigned int posIndex) const {
         // Compute the value of this term in the expansion
         double termVal = 1.0;
-        bool hasDeriv = false;
         if(multiSet_.nzStarts(termInd)==multiSet_.nzStarts(termInd+1)) return 0.; // Value is zero if constant in last dimension
         unsigned int end_idx = multiSet_.nzStarts(termInd+1)-1; // last index in loop
         if(multiSet_.nzDims(end_idx)!=dim_-1) return 0.; // Value is zero if constant in last dimension

MParT/ParameterizedFunctionBase.h

@@ -51,23 +51,21 @@ namespace mpart {
             @param coeffs A view containing the coefficients to copy.
         */
        virtual void SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs);
+       virtual void SetCoeffs(Kokkos::View<double*, MemorySpace> coeffs);
 
-       #if defined(MPART_ENABLE_GPU)
-       void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<MemorySpace,Kokkos::HostSpace>,mpart::DeviceSpace,Kokkos::HostSpace>> coeffs);
-       #endif
 
         /** @brief Wrap the internal coefficient view around another view.
             @details Performs a shallow copy of the input coefficients to the internally stored coefficients.
             If values in the view passed to this function are changed, the values will also change in the
             internally stored view.
             @param coeffs A view containing the coefficients we want to wrap.
         */
-        virtual void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs);
+        virtual void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs);
 
         #if defined(MPART_ENABLE_GPU)
-        virtual void SetCoeffs(Kokkos::View<const double*, mpart::DeviceSpace> coeffs);
-        virtual void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs);
-        #endif
+        virtual void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<Kokkos::HostSpace,MemorySpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs);
+	virtual void SetCoeffs(Kokkos::View<double*, std::conditional_t<std::is_same_v<Kokkos::HostSpace,MemorySpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs);
+	#endif
 
         /** SetCoeffs function with conversion from Eigen to Kokkos vector types.*/
         virtual void SetCoeffs(Eigen::Ref<Eigen::VectorXd> coeffs);
@@ -249,4 +247,4 @@ namespace mpart {
     }; // class ParameterizedFunctionBase
 }
 
-#endif
+#endif

MParT/Sigmoid.h

@@ -109,7 +109,7 @@ class Sigmoid1d {
 						Kokkos::View<double*, MemorySpace> widths)
 			: centers_(centers), widths_(widths) {
 		Kokkos::View<double*, MemorySpace> weights ("Sigmoid weights", centers.extent(0));
-		Kokkos::parallel_for(centers.extent(0), KOKKOS_LAMBDA(unsigned int i) { weights(i) = 1.; });
+		Kokkos::deep_copy(weights, 1.0);
 		weights_ = weights;
 		Validate();
 	}
@@ -121,7 +121,7 @@ class Sigmoid1d {
 	 * @param max_order Maximum order of basis function to evaluate
 	 * @param input Point to evaluate function
 	 */
-	void EvaluateAll(double* output, int max_order, double input) const {
+	KOKKOS_FUNCTION void EvaluateAll(double* output, int max_order, double input) const {
 		if (order_ < max_order) {
 			std::stringstream ss;
 			ss << "Sigmoid basis evaluation order too large.\n";
@@ -164,7 +164,7 @@ class Sigmoid1d {
 	 * @param max_order Number of sigmoids to evaluate
 	 * @param input Where to evaluate sigmoids
 	 */
-	void EvaluateDerivatives(double* output, double* output_diff, int max_order,
+	KOKKOS_FUNCTION void EvaluateDerivatives(double* output, double* output_diff, int max_order,
 							 double input) const {
 		if (order_ < max_order) {
 			std::stringstream ss;
@@ -218,7 +218,7 @@ class Sigmoid1d {
 	 * @param max_order Maximum order of sigmoid to evaluate
 	 * @param input Where to evaluate the sigmoids
 	 */
-	void EvaluateSecondDerivatives(double* output, double* output_diff,
+	KOKKOS_FUNCTION void EvaluateSecondDerivatives(double* output, double* output_diff,
 								   double* output_diff2, int max_order,
 								   double input) const {
 		if (order_ < max_order) {

MParT/SummarizedMap.h

@@ -44,11 +44,10 @@ where the function \f$s:\mathbb{R}^{N-1}\rightarrow \mathbb{R}^{r}\f$ is a funct
     virtual ~SummarizedMap() = default;
 
     using ConditionalMapBase<MemorySpace>::SetCoeffs;
-    void SetCoeffs(Kokkos::View<const double*, Kokkos::HostSpace> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs) override;
+    void SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs) override;
+    void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs) override;
     #if defined(MPART_ENABLE_GPU)
-    void SetCoeffs(Kokkos::View<const double*, DeviceSpace> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, DeviceSpace> coeffs) override;
+    void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<Kokkos::HostSpace, MemorySpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs) override;
     #endif
 
     void SummarizePts(StridedMatrix<const double, MemorySpace> const&  pts,
@@ -95,4 +94,4 @@ where the function \f$s:\mathbb{R}^{N-1}\rightarrow \mathbb{R}^{r}\f$ is a funct
 
 }
 
-#endif
+#endif

MParT/TrainMap.h

@@ -3,8 +3,9 @@
 
 #include <functional>
 #include <nlopt.hpp>
-#include "ConditionalMapBase.h"
-#include "MapObjective.h"
+#include <iostream>
+#include "MParT/ConditionalMapBase.h"
+#include "MParT/MapObjective.h"
 
 namespace mpart {
 

MParT/Utilities/LinearAlgebra.h

@@ -195,55 +195,6 @@ StridedMatrix<double, MemorySpace> operator*(TransposeObject<MemorySpace> A,
     return C;
 }
 
-enum class ReduceDimMap { sum, norm };
-
-template<ReduceDimMap Map>
-KOKKOS_INLINE_FUNCTION double ReduceDimMapKernel(double) {return 0.;};
-
-template<>
-KOKKOS_INLINE_FUNCTION double ReduceDimMapKernel<ReduceDimMap::sum>(double x) {return x;};
-
-template<>
-KOKKOS_INLINE_FUNCTION double ReduceDimMapKernel<ReduceDimMap::norm>(double x) {return x*x;};
-
-
-// This struct allows you to reduce the columns of a mxn matrix
-// Performs alpha*A*[1,...,1]
-template<ReduceDimMap Map, typename MemorySpace, unsigned int Dim = 1, typename = std::enable_if_t<Dim <= 1,int> >
-struct ReduceDim {
-    using value_type = double[];
-    using size_type = typename StridedMatrix<double, MemorySpace>::size_type;
-
-    // Keep track of the dimension we are not reducing
-    size_type value_count;
-
-    StridedMatrix<const double, MemorySpace> A_;
-    double alpha_;
-
-    ReduceDim(StridedMatrix<double, MemorySpace> A, double alpha): value_count(A.extent(1-Dim)), A_(A), alpha_(alpha) {}
-
-    KOKKOS_INLINE_FUNCTION void operator()(const size_type reduce_idx, value_type sum) const {
-        for(size_type full_idx=0; full_idx<value_count; ++full_idx) {
-            double aij = 0.;
-            if(Dim == 1) aij = A_(full_idx,reduce_idx);
-            else aij = A_(reduce_idx,full_idx);
-            sum[full_idx] += ReduceDimMapKernel<Map>(aij)*alpha_;
-        }
-    }
-
-    KOKKOS_INLINE_FUNCTION void join (volatile value_type dst, const volatile value_type src) const {
-        for (size_type i = 0; i < value_count; ++i) {
-            dst[i] += src[i];
-        }
-    }
-
-    KOKKOS_INLINE_FUNCTION void init (value_type sum) const {
-        for (size_type i = 0; i < value_count; ++i) {
-            sum[i] = 0.0;
-        }
-    }
-};
-
 /** Mimics the interface of the Eigen::PartialPivLU class, but using Kokkos::Views and CUBLAS/CUSOLVER linear algebra.
 
 Note that the layout of the matrices used in this class is important.  Cublas expects column major (layout left).

MParT/Utilities/Miscellaneous.h

@@ -24,8 +24,8 @@ namespace mpart{
      */
     template<typename MemorySpace, typename ErrorType>
     struct ProcAgnosticError {
-        static void error(const char*) {
-            assert(false);
+        KOKKOS_INLINE_FUNCTION static void error(const char*) {
+            KOKKOS_ASSERT(false);
         }
     };
 

MParT/Utilities/RootFinding.h

@@ -113,7 +113,6 @@ KOKKOS_INLINE_FUNCTION double Find_x_ITP(double xlb, double xub, double yd, doub
 template<typename MemorySpace, typename FunctorType>
 KOKKOS_INLINE_FUNCTION double InverseSingleBracket(double yd, FunctorType f, double x0, const double xtol, const double ftol, int& info)
 {   
-    double stepSize=1.0;
     const unsigned int maxIts = 10000;
 
     // First, we need to find two points that bound the solution.

bindings/julia/src/CommonJuliaUtilities.cpp

@@ -25,6 +25,5 @@ void mpart::binding::CommonUtilitiesWrapper(jlcxx::Module &mod)
     mod.method("Initialize", [](){mpart::binding::Initialize(std::vector<std::string> {});});
     mod.method("Initialize", [](std::vector<std::string> v){mpart::binding::Initialize(makeInitArguments(v));});
     mod.method("Concurrency", &Kokkos::DefaultExecutionSpace::concurrency);
-    mod.add_type<Kokkos::HostSpace>("HostSpace");
     mod.add_type<Kokkos::LayoutStride>("LayoutStride");
 }

src/ComposedMap.cpp

@@ -178,7 +178,7 @@ ComposedMap<MemorySpace>::ComposedMap(std::vector<std::shared_ptr<ConditionalMap
 }
 
 template<typename MemorySpace>
-void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<const double*, Kokkos::HostSpace> coeffs)
+void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs)
 {
     // First, call the ConditionalMapBase version of this function to copy the view into the savedCoeffs member variable
     ConditionalMapBase<MemorySpace>::SetCoeffs(coeffs);
@@ -195,7 +195,7 @@ void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<const double*, Kokkos::Hos
 }
 
 template<typename MemorySpace>
-void ComposedMap<MemorySpace>::WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs)
+void ComposedMap<MemorySpace>::WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs)
 {
     // First, call the ConditionalMapBase version of this function to copy the view into the savedCoeffs member variable
     ConditionalMapBase<MemorySpace>::WrapCoeffs(coeffs);
@@ -212,19 +212,36 @@ void ComposedMap<MemorySpace>::WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpac
 }
 
 #if defined(MPART_ENABLE_GPU)
-template<typename MemorySpace>
-void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<const double*, Kokkos::DefaultExecutionSpace::memory_space> coeffs)
+template<>
+void ComposedMap<Kokkos::HostSpace>::SetCoeffs(Kokkos::View<const double*, mpart::DeviceSpace> coeffs)
 {
     // First, call the ConditionalMapBase version of this function to copy the view into the savedCoeffs member variable
-    ConditionalMapBase<MemorySpace>::SetCoeffs(coeffs);
+     ConditionalMapBase<Kokkos::HostSpace>::SetCoeffs(coeffs);
 
     // Now create subviews for each of the maps
     unsigned int cumNumCoeffs = 0;
     for(unsigned int i=0; i<maps_.size(); ++i){
         assert(cumNumCoeffs+maps_.at(i)->numCoeffs <= this->savedCoeffs.size());
 
-        maps_.at(i)->savedCoeffs = Kokkos::subview(this->savedCoeffs,
-            std::make_pair(cumNumCoeffs, cumNumCoeffs+maps_.at(i)->numCoeffs));
+        maps_.at(i)->WrapCoeffs(Kokkos::subview(this->savedCoeffs,
+            std::make_pair(cumNumCoeffs, cumNumCoeffs+maps_.at(i)->numCoeffs)));
+        cumNumCoeffs += maps_.at(i)->numCoeffs;
+    }
+}
+
+template<>
+void ComposedMap<mpart::DeviceSpace>::SetCoeffs(Kokkos::View<const double*, Kokkos::HostSpace> coeffs)
+{
+    // First, call the ConditionalMapBase version of this function to copy the view into the savedCoeffs member variable
+     ConditionalMapBase<mpart::DeviceSpace>::SetCoeffs(coeffs);
+
+    // Now create subviews for each of the maps
+    unsigned int cumNumCoeffs = 0;
+    for(unsigned int i=0; i<maps_.size(); ++i){
+        assert(cumNumCoeffs+maps_.at(i)->numCoeffs <= this->savedCoeffs.size());
+
+        maps_.at(i)->WrapCoeffs(Kokkos::subview(this->savedCoeffs,
+            std::make_pair(cumNumCoeffs, cumNumCoeffs+maps_.at(i)->numCoeffs)));
         cumNumCoeffs += maps_.at(i)->numCoeffs;
     }
 }
@@ -471,4 +488,4 @@ void ComposedMap<MemorySpace>::LogDeterminantInputGradImpl(StridedMatrix<const d
 template class mpart::ComposedMap<Kokkos::HostSpace>;
 #if defined(MPART_ENABLE_GPU)
     template class mpart::ComposedMap<DeviceSpace>;
-#endif
+#endif

src/ConditionalMapBase.cpp

@@ -250,12 +250,13 @@ StridedMatrix<double, mpart::DeviceSpace> ConditionalMapBase<mpart::DeviceSpace>
     return output;
 }
 
+
 template<>
 template<>
 StridedMatrix<double, Kokkos::HostSpace> ConditionalMapBase<mpart::DeviceSpace>::LogDeterminantInputGrad(StridedMatrix<const double, Kokkos::HostSpace> const& pts)
 {
     // Copy the points to the device space
-    StridedMatrix<const double, mpart::DeviceSpace> pts_device = ToDevice<mpart::DeviceSpace,const double>(pts);
+    StridedMatrix<const double, mpart::DeviceSpace> pts_device = ToDevice<mpart::DeviceSpace>(pts);
 
     // Evaluate on the device space
     StridedMatrix<double, mpart::DeviceSpace> evals_device = this->LogDeterminantInputGrad(pts_device);
@@ -275,7 +276,7 @@ StridedMatrix<double, mpart::DeviceSpace> ConditionalMapBase<Kokkos::HostSpace>:
     StridedMatrix<double, Kokkos::HostSpace> evals_host = this->LogDeterminantInputGrad(pts_host);
 
     // Copy back to the device
-    return ToDevice<mpart::DeviceSpace,double>(evals_host);
+    return ToDevice<mpart::DeviceSpace>(evals_host);
 }
 
 
@@ -296,4 +297,4 @@ Eigen::RowMatrixXd ConditionalMapBase<mpart::DeviceSpace>::LogDeterminantInputGr
 template class mpart::ConditionalMapBase<Kokkos::HostSpace>;
 #if defined(MPART_ENABLE_GPU)
     template class mpart::ConditionalMapBase<DeviceSpace>;
-#endif
+#endif