1 // RUN: %libomptarget-compilexx-run-and-check-generic
2
3 // Currently hangs on amdgpu
4 // UNSUPPORTED: amdgcn-amd-amdhsa
5 // UNSUPPORTED: amdgcn-amd-amdhsa-oldDriver
6 // UNSUPPORTED: amdgcn-amd-amdhsa-LTO
7 // UNSUPPORTED: x86_64-pc-linux-gnu
8 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
9 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
10
11 #include <cassert>
12 #include <cmath>
13 #include <iostream>
14 #include <limits>
15 #include <memory>
16 #include <vector>
17
18 class BlockMatrix {
19 private:
20 const int rowsPerBlock;
21 const int colsPerBlock;
22 const long nRows;
23 const long nCols;
24 const int nBlocksPerRow;
25 const int nBlocksPerCol;
26 std::vector<std::vector<std::unique_ptr<float[]>>> Blocks;
27
28 public:
BlockMatrix(const int _rowsPerBlock,const int _colsPerBlock,const long _nRows,const long _nCols)29 BlockMatrix(const int _rowsPerBlock, const int _colsPerBlock,
30 const long _nRows, const long _nCols)
31 : rowsPerBlock(_rowsPerBlock), colsPerBlock(_colsPerBlock), nRows(_nRows),
32 nCols(_nCols), nBlocksPerRow(_nRows / _rowsPerBlock),
33 nBlocksPerCol(_nCols / _colsPerBlock), Blocks(nBlocksPerCol) {
34 for (int i = 0; i < nBlocksPerCol; i++) {
35 for (int j = 0; j < nBlocksPerRow; j++) {
36 Blocks[i].emplace_back(new float[_rowsPerBlock * _colsPerBlock]);
37 }
38 }
39 };
40
41 // Initialize the BlockMatrix from 2D arrays
Initialize(const std::vector<float> & matrix)42 void Initialize(const std::vector<float> &matrix) {
43 for (int i = 0; i < nBlocksPerCol; i++)
44 for (int j = 0; j < nBlocksPerRow; j++) {
45 float *CurrBlock = GetBlock(i, j);
46 for (int ii = 0; ii < colsPerBlock; ++ii)
47 for (int jj = 0; jj < rowsPerBlock; ++jj) {
48 int curri = i * colsPerBlock + ii;
49 int currj = j * rowsPerBlock + jj;
50 CurrBlock[ii + jj * colsPerBlock] = matrix[curri + currj * nCols];
51 }
52 }
53 }
54
Compare(const std::vector<float> & matrix) const55 void Compare(const std::vector<float> &matrix) const {
56 for (int i = 0; i < nBlocksPerCol; i++)
57 for (int j = 0; j < nBlocksPerRow; j++) {
58 float *CurrBlock = GetBlock(i, j);
59 for (int ii = 0; ii < colsPerBlock; ++ii)
60 for (int jj = 0; jj < rowsPerBlock; ++jj) {
61 int curri = i * colsPerBlock + ii;
62 int currj = j * rowsPerBlock + jj;
63 float m_value = matrix[curri + currj * nCols];
64 float bm_value = CurrBlock[ii + jj * colsPerBlock];
65 assert(std::fabs(bm_value - m_value) <
66 std::numeric_limits<float>::epsilon());
67 }
68 }
69 }
70
GetBlock(int i,int j) const71 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
BlockMatMul_TargetNowait(BlockMatrix & A,BlockMatrix & B,BlockMatrix & C)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
Matmul(const std::vector<float> & a,const std::vector<float> & b,std::vector<float> & c)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
main(int argc,char * argv[])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 auto BlockedB = BlockMatrix(BS, BS, N, N);
132 auto BlockedC = BlockMatrix(BS, BS, N, N);
133 BlockedA.Initialize(a);
134 BlockedB.Initialize(b);
135 BlockedC.Initialize(c);
136 BlockedA.Compare(a);
137 BlockedB.Compare(b);
138 BlockedC.Compare(c);
139
140 Matmul(a, b, c);
141 BlockMatMul_TargetNowait(BlockedA, BlockedB, BlockedC);
142
143 BlockedC.Compare(c);
144
145 std::cout << "PASS\n";
146
147 return 0;
148 }
149
150 // CHECK: PASS
151