| 109 | } |
| 110 | |
| 111 | Kernel selectTranspose(Context &ctx, int version, |
| 112 | const Bindings</* input, output */ 2> &bindings, |
| 113 | size_t M, size_t N) { |
| 114 | Kernel kernel; |
| 115 | if (version == 1) { |
| 116 | Shape wgSize = {16, 16, 1}; |
| 117 | LOG(kDefLog, kInfo, "wgSize: %s", toString(wgSize).c_str()); |
| 118 | KernelCode transpose = |
| 119 | createTranspose1(kShaderTranspose1, M, N, /*wgsize*/ wgSize); // The shape of input == M x N |
| 120 | kernel = createKernel(ctx, transpose, bindings, |
| 121 | /*nWorkgroups*/ cdiv({N, M, 1}, wgSize)); // The shape of output == N x M |
| 122 | } else if (version == 2) { |
| 123 | static constexpr size_t BM = 64; |
| 124 | static constexpr size_t BK = 16; |
| 125 | static constexpr size_t BN = 64; |
| 126 | static constexpr size_t TM = BM / BK; |
| 127 | static constexpr size_t TN = BN / BK; |
| 128 | Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. |
| 129 | Shape nWorkgroups = {cdiv(N, BN), cdiv(M, BM), 1}; |
| 130 | LOG(kDefLog, kInfo, "M: %d, N: %d", M, N); |
| 131 | LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); |
| 132 | LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); |
| 133 | LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); |
| 134 | KernelCode transpose = createTranspose2(kShaderTranspose2, M, N, BM, BN, TM, TN, |
| 135 | /*wgSize*/ wgSize, |
| 136 | kf32); |
| 137 | kernel = createKernel(ctx, transpose, bindings, |
| 138 | /*nWorkgroups*/ nWorkgroups); |
| 139 | } else if (version == 0) { |
| 140 | LOG(kDefLog, kInfo, "Skip Creating Kernel", M, N); |
| 141 | } |
| 142 | return kernel; |
| 143 | } |
| 144 | |
| 145 | void runTest(int version, size_t M, size_t N, |
| 146 | std::unique_ptr<float[]> &inputPtr, |
no test coverage detected