1.04 - MDSpan

Exercise - MDSpan

Input

%%writefile Sources/heat-2D.cu
#include "dli.h"

__host__ __device__
cuda::std::pair<int, int> row_col(int id, int width) {
    return cuda::std::make_pair(id / width, id % width);
}

void simulate(int height, int width,
              const thrust::universal_vector<float> &in,
                    thrust::universal_vector<float> &out)
{
  const float *in_ptr = thrust::raw_pointer_cast(in.data());

  thrust::tabulate(
    thrust::device, out.begin(), out.end(), 
    [in_ptr, height, width] __host__ __device__(int id) {
      auto [row, column] = row_col(id, width);

      if (row > 0 && column > 0 && row < height - 1 && column < width - 1) {
        float d2tdx2 = in_ptr[(row) * width + column - 1] - 2 * in_ptr[row * width + column] + in_ptr[(row) * width + column + 1];
        float d2tdy2 = in_ptr[(row - 1) * width + column] - 2 * in_ptr[row * width + column] + in_ptr[(row + 1) * width + column];

        return in_ptr[row * width + column] + 0.2f * (d2tdx2 + d2tdy2);
      } else {
        return in_ptr[row * width + column];
      }
    });
}

Output

mdspan needs to utilize an input, which needs to be a raw pointer to some sort of iterable dataset. We're given in_data() in_ptr which we can use with height and width of the actual underlying grid to retrieve values. mdspan just needs height and width to do the correct modulo and division operations to efficiently index the right location

mdspan is similar to iloc in pandas / numpy on python - helps to get retrieve $N$-dimensional coordinates even though the array in memory is just a flat structure

%%writefile Sources/heat-2D.cu
#include "dli.h"

__host__ __device__
cuda::std::pair<int, int> row_col(int id, int width) {
    return cuda::std::make_pair(id / width, id % width);
}

void simulate(int height, int width,
              const thrust::universal_vector<float> &in,
                    thrust::universal_vector<float> &out)
{
  const float *in_ptr = thrust::raw_pointer_cast(in.data());

  thrust::tabulate(
    thrust::device, out.begin(), out.end(), 
    [in_ptr, height, width] __host__ __device__(int id) {
      auto [row, col] = row_col(id, width);
      cuda::std::mdspan md(in_ptr, height, width);

      if (row > 0 && col > 0 && row < height - 1 && col < width - 1) {
        float d2tdx2 = md(row, col - 1) - 2 * md(row, col) + md(row, col + 1);
        float d2tdy2 = md(row - 1, col) - 2 * md(row, col) + md(row + 1, col);

        return md(row, col) + 0.2f * (d2tdx2 + d2tdy2);
      } else {
        return md(row, col);
      }
    });
}