                             cudaMemcpyDeviceToHost));
 }
 
-__global__ void Pack3D(const int *src, int *dst,
+//! Pack a 3-dimensional sub grid into a linear region
+/*!
+  \param src 3D source grid 
+  \param dst Linear buffer
+  \param xdim X size of the grid
+  \param ydim Y size of the grid
+  \param zdim Z size of the grid
+  \param xoff X offset of the sub grid
+  \param yoff Y offset of the sub grid
+  \param zoff Z offset of the sub grid
+  \param xextent X size of the sub grid
+  \param yextent Y size of the sub grid
+  \param zextent Z size of the sub grid 
+ */
+template <class T>
+__global__ void Pack3D(const T *src, T *dst,
                        PSIndex xdim, PSIndex ydim, PSIndex zdim,
                        PSIndex xoff, PSIndex yoff, PSIndex zoff,
                        PSIndex xextent, PSIndex yextent, PSIndex zextent) {

    CUDA_SAFE_CALL(cudaMemcpy(temp_buffer_->Get(), subgrid, 
                              subgrid_size.accumulate(rank) * elm_size,
                              cudaMemcpyHostToDevice));
    
    if (rank == 3) {
      Unpack3d<<<nblocks, thread_block_size>>>(
          (int*)(temp_buffer_->Get()), (int*)Get(),
          grid_size[0], grid_size[1], grid_size[2],
          subgrid_offset[0], subgrid_offset[1], subgrid_offset[2],
          subgrid_size[0], subgrid_size[1], subgrid_size[2]);
    } else if (rank == 2) {
      Unpack2d<<<nblocks, thread_block_size>>>(
          (int*)(temp_buffer_->Get()), (int*)Get(), 
          grid_size[0], grid_size[1], 
          subgrid_offset[0], subgrid_offset[1],
          subgrid_size[0], subgrid_size[1]);
    } else {
      LOG_ERROR() << "Unsupported\n";
      PSAbort(1);

                            cudaMemcpyHostToDevice));
}

template <class T>
__global__ void Unpack2d(const T *src, T *dst,
                         PSIndex xdim, PSIndex ydim,
                         PSIndex xoff, PSIndex yoff,
                         PSIndex xextent, PSIndex yextent) {
