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