how to cast a 2-dimensional thrust::device_vector<thrust::device_vector<int>> to raw pointer

Question

When I use the thrust::device_vector in main function,I can pass it to the kernel function correctly,the code as follows:

 thrust::device_vector<int> device_a(2);
 thrust::host_vector<int> host_a(2);
 MyTest << <1, 2 >> >(thrust::raw_pointer_cast(&device_a[0]),device_a.size());
 host_a = device_a;
 for (int i = 0; i < host_a.size();i++)
 cout << host_a[i] << endl;

but I want to use 2-dimension device_vector in my code,How can I use it? as shown i the following code

__global__ void MyTest(thrust::device_vector<int>* a,int total){  
    int idx = threadIdx.x;  
    if (idx < total){
        int temp = idx;  
        a[idx][0] = temp;  
        a[idx][1] = temp; 
        __syncthreads();
      }

}  
 void main(){
    thrust::device_vector<thrust::device_vector<int>> device_a(2,thrust::device_vector<int>(2));

    thrust::host_vector<thrust::host_vector<int>> host_a(2,thrust::host_vector<int>(2));

    MyTest << <1, 2 >> >(thrust::raw_pointer_cast(device_a.data()),device_a.size());
    host_a = device_a;
    for (int i = 0; i < host_a.size(); i++){
    cout << host_a[i][0] << endl;
    cout << host_a[i][1] << endl;
}
}

Answer 1

Generally, Thrust containers are host only types that can not be used in __device__ and __global__ functions.

The common way to use 2-D array is to put it in a 1-D linear memory space like the following code.

__global__ void MyTest(int* a, int nrows, int ncols) {
  int j = threadIdx.x;
  int i = threadIdx.y;
  if (i < nrows && j < ncols) {
    int temp = i + j;
    a[i * ncols + j] = temp;
  }

}

int main(int argc, char** argv) {
  int nrows = 2;
  int ncols = 2;
  thrust::device_vector<int> device_a(nrows * ncols);
  MyTest<<<1, dim3(2, 2)>>>(thrust::raw_pointer_cast(device_a.data()), rows, ncols);
  return 0;
}