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