offload/test/offloading/bug49334.cpp

   1 // RUN: %libomptarget-compilexx-generic -O3 && %libomptarget-run-generic
   2 // RUN: %libomptarget-compilexx-generic -O3 -ffast-math && \
   3 // RUN: %libomptarget-run-generic
   4 // RUN: %libomptarget-compileoptxx-generic -O3 && %libomptarget-run-generic
   5 // RUN: %libomptarget-compileoptxx-generic -O3 -ffast-math && \
   6 // RUN: %libomptarget-run-generic
   7
   8 // REQUIRES: gpu
   9 // UNSUPPORTED: nvidiagpu
  10 // UNSUPPORTED: amdgpu
  11
  12 #include <cassert>
  13 #include <cmath>
  14 #include <iostream>
  15 #include <limits>
  16 #include <memory>
  17 #include <vector>
  18
  19 class BlockMatrix {
  20 private:
  21   const int rowsPerBlock;
  22   const int colsPerBlock;
  23   const long nRows;
  24   const long nCols;
  25   const int nBlocksPerRow;
  26   const int nBlocksPerCol;
  27   std::vector<std::vector<std::unique_ptr<float[]>>> Blocks;
  28
  29 public:
  30   BlockMatrix(const int _rowsPerBlock, const int _colsPerBlock,
  31               const long _nRows, const long _nCols)
  32       : rowsPerBlock(_rowsPerBlock), colsPerBlock(_colsPerBlock), nRows(_nRows),
  33         nCols(_nCols), nBlocksPerRow(_nRows / _rowsPerBlock),
  34         nBlocksPerCol(_nCols / _colsPerBlock), Blocks(nBlocksPerCol) {
  35     for (int i = 0; i < nBlocksPerCol; i++) {
  36       for (int j = 0; j < nBlocksPerRow; j++) {
  37         Blocks[i].emplace_back(new float[_rowsPerBlock * _colsPerBlock]);
  38       }
  39     }
  40   };
  41
  42   // Initialize the BlockMatrix from 2D arrays
  43   void Initialize(const std::vector<float> &matrix) {
  44     for (int i = 0; i < nBlocksPerCol; i++)
  45       for (int j = 0; j < nBlocksPerRow; j++) {
  46         float *CurrBlock = GetBlock(i, j);
  47         for (int ii = 0; ii < colsPerBlock; ++ii)
  48           for (int jj = 0; jj < rowsPerBlock; ++jj) {
  49             int curri = i * colsPerBlock + ii;
  50             int currj = j * rowsPerBlock + jj;
  51             CurrBlock[ii + jj * colsPerBlock] = matrix[curri + currj * nCols];
  52           }
  53       }
  54   }
  55
  56   void Compare(const std::vector<float> &matrix) const {
  57     for (int i = 0; i < nBlocksPerCol; i++)
  58       for (int j = 0; j < nBlocksPerRow; j++) {
  59         float *CurrBlock = GetBlock(i, j);
  60         for (int ii = 0; ii < colsPerBlock; ++ii)
  61           for (int jj = 0; jj < rowsPerBlock; ++jj) {
  62             int curri = i * colsPerBlock + ii;
  63             int currj = j * rowsPerBlock + jj;
  64             float m_value = matrix[curri + currj * nCols];
  65             float bm_value = CurrBlock[ii + jj * colsPerBlock];
  66             assert(std::fabs(bm_value - m_value) <
  67                    std::numeric_limits<float>::epsilon());
  68           }
  69       }
  70   }
  71
  72   float *GetBlock(int i, int j) const {
  73     assert(i < nBlocksPerCol && j < nBlocksPerRow && "Accessing outside block");
  74     return Blocks[i][j].get();
  75   }
  76 };
  77
  78 constexpr const int BS = 16;
  79 constexpr const int N = 256;
  80
  81 int BlockMatMul_TargetNowait(BlockMatrix &A, BlockMatrix &B, BlockMatrix &C) {
  82 #pragma omp parallel
  83 #pragma omp master
  84   for (int i = 0; i < N / BS; ++i)
  85     for (int j = 0; j < N / BS; ++j) {
  86       float *BlockC = C.GetBlock(i, j);
  87       for (int k = 0; k < N / BS; ++k) {
  88         float *BlockA = A.GetBlock(i, k);
  89         float *BlockB = B.GetBlock(k, j);
  90 // clang-format off
  91 #pragma omp target depend(in: BlockA[0], BlockB[0]) depend(inout: BlockC[0])   \
  92             map(to: BlockA[:BS * BS], BlockB[:BS * BS])                        \
  93             map(tofrom: BlockC[:BS * BS]) nowait
  94 // clang-format on
  95 #pragma omp parallel for
  96         for (int ii = 0; ii < BS; ii++)
  97           for (int jj = 0; jj < BS; jj++) {
  98             for (int kk = 0; kk < BS; ++kk)
  99               BlockC[ii + jj * BS] +=
 100                   BlockA[ii + kk * BS] * BlockB[kk + jj * BS];
 101           }
 102       }
 103     }
 104   return 0;
 105 }
 106
 107 void Matmul(const std::vector<float> &a, const std::vector<float> &b,
 108             std::vector<float> &c) {
 109   for (int i = 0; i < N; ++i) {
 110     for (int j = 0; j < N; ++j) {
 111       float sum = 0.0;
 112       for (int k = 0; k < N; ++k) {
 113         sum = sum + a[i * N + k] * b[k * N + j];
 114       }
 115       c[i * N + j] = sum;
 116     }
 117   }
 118 }
 119
 120 int main(int argc, char *argv[]) {
 121   std::vector<float> a(N * N);
 122   std::vector<float> b(N * N);
 123   std::vector<float> c(N * N, 0.0);
 124
 125   for (int i = 0; i < N; ++i) {
 126     for (int j = 0; j < N; ++j) {
 127       a[i * N + j] = b[i * N + j] = i + j % 100;
 128     }
 129   }
 130
 131   auto BlockedA = BlockMatrix(BS, BS, N, N);
 132   auto BlockedB = BlockMatrix(BS, BS, N, N);
 133   auto BlockedC = BlockMatrix(BS, BS, N, N);
 134   BlockedA.Initialize(a);
 135   BlockedB.Initialize(b);
 136   BlockedC.Initialize(c);
 137   BlockedA.Compare(a);
 138   BlockedB.Compare(b);
 139   BlockedC.Compare(c);
 140
 141   Matmul(a, b, c);
 142   BlockMatMul_TargetNowait(BlockedA, BlockedB, BlockedC);
 143
 144   BlockedC.Compare(c);
 145
 146   std::cout << "PASS\n";
 147
 148   return 0;
 149 }
 150
 151 // CHECK: PASS