1 files changed, 61 insertions, 1 deletions
diff --git a/chroma/cuda/propagate.cu b/chroma/cuda/propagate.cu
index fdcf532..87ba9f5 100644
--- a/chroma/cuda/propagate.cu
+++ b/chroma/cuda/propagate.cu
@@ -45,7 +45,67 @@ photon_duplicate(int first_photon, int nthreads,
       histories[target_photon_id] = p.history;
     }
 }
-		 
+
+__global__ void
+count_photons(int first_photon, int nthreads, unsigned int target_flag,
+	      unsigned int *index_counter,
+	      unsigned int *histories)
+{
+    int id = blockIdx.x*blockDim.x + threadIdx.x;
+    __shared__ unsigned int counter;
+
+    if (threadIdx.x == 0)
+	counter = 0;
+    __syncthreads();
+
+    if (id < nthreads) {
+	int photon_id = first_photon + id;
+
+	if (histories[photon_id] & target_flag) {
+	    atomicAdd(&counter, 1);
+	}
+	    
+    }
+
+    __syncthreads();
+
+    if (threadIdx.x == 0)
+	atomicAdd(index_counter, counter);
+}
+
+__global__ void
+copy_photons(int first_photon, int nthreads, unsigned int target_flag,
+	     unsigned int *index_counter,
+	     float3 *positions, float3 *directions,
+	     float *wavelengths, float3 *polarizations,
+	     float *times, unsigned int *histories,
+	     int *last_hit_triangles, 
+	     float3 *new_positions, float3 *new_directions,
+	     float *new_wavelengths, float3 *new_polarizations,
+	     float *new_times, unsigned int *new_histories,
+	     int *new_last_hit_triangles)
+{
+    int id = blockIdx.x*blockDim.x + threadIdx.x;
+    
+    if (id >= nthreads)
+	return;
+    
+    int photon_id = first_photon + id;
+
+    if (histories[photon_id] & target_flag) {
+	int offset = atomicAdd(index_counter, 1);
+
+	new_positions[offset] = positions[photon_id];
+	new_directions[offset] = directions[photon_id];
+	new_polarizations[offset] = polarizations[photon_id];
+	new_wavelengths[offset] = wavelengths[photon_id];
+	new_times[offset] = times[photon_id];
+	new_histories[offset] = histories[photon_id];
+	new_last_hit_triangles[offset] = last_hit_triangles[photon_id];
+    }
+}
+
+	      
 __global__ void
 propagate(int first_photon, int nthreads, unsigned int *input_queue,
 	  unsigned int *output_queue, curandState *rng_states,