Merge pull request #27 from zhihuadong/full1

Move response calculation to device.
WireCell · Apr 14, 2022 · 06f0d60 · 06f0d60
2 parents 997d1e9 + 0c45533
commit 06f0d60
Show file tree

Hide file tree

Showing 11 changed files with 388 additions and 88 deletions.
diff --git a/.cmake-kk-hip/CMakeLists.txt b/.cmake-kk-hip/CMakeLists.txt
@@ -36,6 +36,7 @@ add_library(WireCellGenKokkos SHARED ${all_files})
 target_include_directories(WireCellGenKokkos
   PRIVATE
     ${PROJECT_SOURCE_DIR}/../inc
+    $ENV{HIPFFT_INC}
     $ENV{EIGEN_INC}
     $ENV{JSONCPP_INC}
     $ENV{JSONNET_INC}

diff --git a/.cmake-kokkos-omp/CMakeLists.txt b/.cmake-kokkos-omp/CMakeLists.txt
@@ -36,6 +36,7 @@ add_library(WireCellGenKokkos SHARED ${all_files})
 target_include_directories(WireCellGenKokkos
   PRIVATE
     ${PROJECT_SOURCE_DIR}/../inc
+    $ENV{FFTW_INC}
     $ENV{EIGEN_INC}
     $ENV{JSONCPP_INC}
     $ENV{JSONNET_INC}

diff --git a/example/sim.fcl b/example/sim.fcl
@@ -1,7 +1,8 @@
 process_name: wclssim
 
 # "Pgrapher" or "TbbFlow"
-wc_engine: "Pgrapher"
+#wc_engine: "Pgrapher"
+wc_engine: "TbbFlow"
 
 physics :{
    producers: {

diff --git a/example/sim.jsonnet b/example/sim.jsonnet
@@ -1,10 +1,10 @@
 // Energy depo -> detector response simulation -> waveform frames
 
 // local reality = std.extVar('reality');
-// local engine = std.extVar('engine');
+local engine = std.extVar('engine');
 
 local reality = 'data';
-local engine = 'Pgrapher';
+//local engine = 'Pgrapher';
 
 local g = import 'pgraph.jsonnet';
 local f = import 'pgrapher/common/funcs.jsonnet';

diff --git a/inc/WireCellGenKokkos/BinnedDiffusion_transform.h b/inc/WireCellGenKokkos/BinnedDiffusion_transform.h
@@ -166,7 +166,7 @@ namespace WireCell {
             //Kokkos::DualView<double[MAX_NPSS_DEVICE]> m_pvec;
             //Kokkos::DualView<double[MAX_NTSS_DEVICE]> m_tvec;
 //            Kokkos::View<float*> m_patch;
-            Kokkos::View<double*> m_normals;
+//            Kokkos::View<double*> m_normals;
 //            Kokkos::DualView<double*> m_ptvecs;
 	    //for batch
 //            Kokkos::View<double*> m_pvecs;

diff --git a/inc/WireCellGenKokkos/KokkosArray.h b/inc/WireCellGenKokkos/KokkosArray.h
@@ -53,7 +53,9 @@ namespace WireCell {
         template <class ViewType>
         inline ViewType Zero(const Index N0, const Index N1)
         {
-            return gen_2d_view<ViewType>(N0, N1, 0);
+	    ViewType ret("ret", N0, N1) ;
+            //return gen_2d_view<ViewType>(N0, N1, 0);
+	    return ret ;
         }
 
         /// Dump out a string for pinting for a 2D view.

diff --git a/inc/WireCellGenKokkos/KokkosArray_cuda.h b/inc/WireCellGenKokkos/KokkosArray_cuda.h
@@ -64,6 +64,33 @@ namespace WireCell {
 
             return out;
         }
+        inline  void dft_rc(const array_xxf& in, array_xxc& out, int dim = 0)
+        {
+            std::cout << "WIRECELL_KOKKOSARRAY_CUDA" << std::endl;
+            Index N0 = in.extent(0);
+            Index N1 = in.extent(1);
+
+            cufftHandle plan;
+
+            if (dim == 0) {
+                int n[] = {(int) N1};
+                int inembed[] = {(int) N1};
+                int onembed[] = {(int) N1};
+                cufftPlanMany(&plan, 1, n, inembed, (int) N0, 1, onembed, (int) N0, 1, CUFFT_R2C, (int) N0);
+                cufftExecR2C(plan, (cufftReal*) in.data(), (cufftComplex*) out.data());
+                cufftDestroy(plan);
+            }
+
+            if (dim == 1) {
+                int n[] = {(int) N0};
+                int inembed[] = {(int) N0};
+                int onembed[] = {(int) N0};
+                cufftPlanMany(&plan, 1, n, inembed, 1, (int) N0, onembed, 1, (int) N0, CUFFT_R2C, (int) N1);
+                cufftExecR2C(plan, (cufftReal*) in.data(), (cufftComplex*) out.data());
+                cufftDestroy(plan);
+            }
+
+        }
         inline array_xxc dft_cc(const array_xxc& in, int dim = 0)
         {
             Index N0 = in.extent(0);
@@ -92,6 +119,31 @@ namespace WireCell {
 
             return out;
         }
+        inline void dft_cc(const array_xxc& in, array_xxc& out, int dim = 0)
+        {
+            Index N0 = in.extent(0);
+            Index N1 = in.extent(1);
+
+            cufftHandle plan;
+
+            if (dim == 0) {
+                int n[] = {(int) N1};
+                int inembed[] = {(int) N1};
+                int onembed[] = {(int) N1};
+                cufftPlanMany(&plan, 1, n, inembed, (int) N0, 1, onembed, (int) N0, 1, CUFFT_C2C, (int) N0);
+                cufftExecC2C(plan, (cufftComplex*) in.data(), (cufftComplex*) out.data(), CUFFT_FORWARD);
+                cufftDestroy(plan);
+            }
+
+            if (dim == 1) {
+                int n[] = {(int) N0};
+                int inembed[] = {(int) N0};
+                int onembed[] = {(int) N0};
+                cufftPlanMany(&plan, 1, n, inembed, 1, (int) N0, onembed, 1, (int) N0, CUFFT_C2C, (int) N1);
+                cufftExecC2C(plan, (cufftComplex*) in.data(), (cufftComplex*) out.data(), CUFFT_FORWARD);
+                cufftDestroy(plan);
+            }
+        }
         inline array_xxc idft_cc(const array_xxc& in, int dim = 0)
         {
             Index N0 = in.extent(0);
@@ -126,6 +178,38 @@ namespace WireCell {
 
             return out;
         }
+        inline void  idft_cc(const array_xxc& in,  array_xxc& out, int dim = 0)
+        {
+            Index N0 = in.extent(0);
+            Index N1 = in.extent(1);
+
+            cufftHandle plan;
+
+            if (dim == 0) {
+                int n[] = {(int) N1};
+                int inembed[] = {(int) N1};
+                int onembed[] = {(int) N1};
+                cufftPlanMany(&plan, 1, n, inembed, (int) N0, 1, onembed, (int) N0, 1, CUFFT_C2C, (int) N0);
+                cufftExecC2C(plan, (cufftComplex*) in.data(), (cufftComplex*) out.data(), CUFFT_INVERSE);
+                cufftDestroy(plan);
+                Kokkos::parallel_for(
+                    Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>({0, 0}, {N0, N1}),
+                    KOKKOS_LAMBDA(const KokkosArray::Index& i0, const KokkosArray::Index& i1) { out(i0, i1) /= N1; });
+            }
+
+            if (dim == 1) {
+                int n[] = {(int) N0};
+                int inembed[] = {(int) N0};
+                int onembed[] = {(int) N0};
+                cufftPlanMany(&plan, 1, n, inembed, 1, (int) N0, onembed, 1, (int) N0, CUFFT_C2C, (int) N1);
+                cufftExecC2C(plan, (cufftComplex*) in.data(), (cufftComplex*) out.data(), CUFFT_INVERSE);
+                cufftDestroy(plan);
+                Kokkos::parallel_for(
+                    Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>({0, 0}, {N0, N1}),
+                    KOKKOS_LAMBDA(const KokkosArray::Index& i0, const KokkosArray::Index& i1) { out(i0, i1) /= N0; });
+            }
+        }
+
         inline array_xxf idft_cr(const array_xxc& in, int dim = 0)
         {
             Index N0 = in.extent(0);

diff --git a/inc/WireCellGenKokkos/KokkosArray_hip.h b/inc/WireCellGenKokkos/KokkosArray_hip.h
@@ -191,6 +191,71 @@ namespace WireCell {
 
             return out;
         }
+        inline void dft_cc(const array_xxc& in, array_xxc& out, int dim = 0)
+        {
+            Index N0 = in.extent(0);
+            Index N1 = in.extent(1);
+
+            hipfftHandle plan ;
+
+            hipfftResult status ;
+            hipError_t sync_status ;
+
+	    if (dim == 0) {
+                int n[] = { (int) N1};
+                int inembed[] = {(int) N1};
+                int onembed[] = {(int) N1};
+                size_t worksize = 0;
+                status = hipfftCreate(&plan);
+                assert(status == HIPFFT_SUCCESS) ;	
+	        // set to autoAllocate
+	        status = hipfftSetAutoAllocation( plan, 1);
+        	assert(status == HIPFFT_SUCCESS) ;
+
+		//MakePlan
+	        status = hipfftMakePlanMany( plan, 1, n, inembed, (int) N0, 1, onembed, N0, 1, HIPFFT_C2C, (int) N0, &worksize);
+        	//std::cout<<"worksize= "<<worksize << std::endl ;
+        	assert(status == HIPFFT_SUCCESS) ;
+
+		//Excute
+		status = hipfftExecC2C( plan,  (hipfftComplex * )in.data(), (hipfftComplex *) out.data() , HIPFFT_FORWARD) ;
+		assert(status == HIPFFT_SUCCESS) ;
+
+		// Wait for execution to finish
+		sync_status=hipDeviceSynchronize() ;
+
+		//Destroy plan
+		hipfftDestroy(plan);
+            }
+
+	    if (dim == 1) {
+                int n[] = {(int) N0};
+		int inembed[] = {(int) N0};
+                int onembed[] = {(int) N0};
+                size_t worksize = 0;
+                status = hipfftCreate(&plan);
+                assert(status == HIPFFT_SUCCESS) ;
+                // set to autoAllocate
+                status = hipfftSetAutoAllocation( plan, 1);
+                assert(status == HIPFFT_SUCCESS) ;
+
+                //MakePlan
+                status = hipfftMakePlanMany( plan, 1, n, inembed,  1, N0, onembed, 1, N0, HIPFFT_C2C, (int) N1, &worksize);
+                //std::cout<<"worksize= "<<worksize << std::endl ;
+                assert(status == HIPFFT_SUCCESS) ;
+
+                //Excute
+                status = hipfftExecC2C( plan,  (hipfftComplex  * )in.data(), (hipfftComplex *) out.data() , HIPFFT_FORWARD) ;
+                assert(status == HIPFFT_SUCCESS) ;
+
+                // Wait for execution to finish
+		sync_status=hipDeviceSynchronize() ;
+
+                //Destroy plan
+                hipfftDestroy(plan);
+
+            }
+        }
         inline array_xxc idft_cc(const array_xxc& in, int dim = 0)
         {
             Index N0 = in.extent(0);
@@ -265,6 +330,77 @@ namespace WireCell {
 
             return out;
         }
+        inline void idft_cc(const array_xxc& in, array_xxc& out , int dim = 0)
+        {
+            Index N0 = in.extent(0);
+            Index N1 = in.extent(1);
+
+            hipfftHandle plan ;
+
+            hipfftResult status ;
+            hipError_t sync_status ;
+
+	    if (dim == 0) {
+                int n[] = { (int) N1};
+                int inembed[] = {(int) N1};
+                int onembed[] = {(int) N1};
+                size_t worksize = 0;
+                status = hipfftCreate(&plan);
+                assert(status == HIPFFT_SUCCESS) ;	
+	        // set to autoAllocate
+	        status = hipfftSetAutoAllocation( plan, 1);
+        	assert(status == HIPFFT_SUCCESS) ;
+
+		//MakePlan
+	        status = hipfftMakePlanMany( plan, 1, n, inembed, (int) N0, 1, onembed, N0, 1, HIPFFT_C2C, (int) N0, &worksize);
+        	//std::cout<<"worksize= "<<worksize << std::endl ;
+        	assert(status == HIPFFT_SUCCESS) ;
+
+		//Excute
+		status = hipfftExecC2C( plan,  (float2 * )in.data(), (float2 *) out.data(), HIPFFT_BACKWARD ) ;
+		assert(status == HIPFFT_SUCCESS) ;
+
+		// Wait for execution to finish
+		sync_status=hipDeviceSynchronize() ;
+
+		//Destroy plan
+		hipfftDestroy(plan);
+		Kokkos::parallel_for(
+                    Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>({0, 0}, {N0, N1}),
+                    KOKKOS_LAMBDA(const KokkosArray::Index& i0, const KokkosArray::Index& i1) { out(i0, i1) /= N1; });
+            }
+
+	    if (dim == 1) {
+                int n[] = {(int) N0};
+		int inembed[] = {(int) N0};
+                int onembed[] = {(int) N0};
+                size_t worksize = 0;
+                status = hipfftCreate(&plan);
+                assert(status == HIPFFT_SUCCESS) ;
+                // set to autoAllocate
+                status = hipfftSetAutoAllocation( plan, 1);
+                assert(status == HIPFFT_SUCCESS) ;
+
+                //MakePlan
+                status = hipfftMakePlanMany( plan, 1, n, inembed, 1 ,N0, onembed, 1, N0, HIPFFT_C2C, (int) N1, &worksize);
+                //std::cout<<"worksize= "<<worksize << std::endl ;
+                assert(status == HIPFFT_SUCCESS) ;
+
+                //Excute
+                status = hipfftExecC2C( plan,  (float2 * )in.data(), (float2 *) out.data() ,HIPFFT_BACKWARD) ;
+                assert(status == HIPFFT_SUCCESS) ;
+
+                // Wait for execution to finish
+		sync_status=hipDeviceSynchronize() ;
+
+                //Destroy plan
+                hipfftDestroy(plan);
+		Kokkos::parallel_for(
+                    Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>({0, 0}, {N0, N1}),
+                    KOKKOS_LAMBDA(const KokkosArray::Index& i0, const KokkosArray::Index& i1) { out(i0, i1) /= N0; });
+
+            }
+        }
         inline array_xxf idft_cr(const array_xxc& in, int dim = 0)
         {
             Index N0 = in.extent(0);