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