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