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