| 457 | )"; |
| 458 | |
| 459 | inline KernelCode createMatmulWithVectorization(const char *shaderTemplate, const size_t M, |
| 460 | const size_t K, const size_t N, const size_t BM, |
| 461 | const size_t BK, const size_t BN, |
| 462 | const size_t TM, const size_t TN, |
| 463 | const Shape &workgroupSize = {256, 1, 1}, |
| 464 | NumType precision = kf32, |
| 465 | bool unrolling = false) { |
| 466 | assert(BM % TM == 0); |
| 467 | assert(BN % TN == 0); |
| 468 | assert(K % BK == 0); |
| 469 | assert(M % BM == 0); |
| 470 | assert(N % BN == 0); |
| 471 | // # threads = tile A size == tile B size == # threads for computing C |
| 472 | int num_threads = BM * BN / (TM * TN); |
| 473 | std::string codeString(shaderTemplate); |
| 474 | replaceAll(codeString, {{"{{workgroupSize}}", toString(workgroupSize)}, |
| 475 | {"{{precision}}", toString(precision)}, |
| 476 | {"{{M}}", toString(M)}, |
| 477 | {"{{K}}", toString(K)}, |
| 478 | {"{{N}}", toString(N)}, |
| 479 | {"{{BM}}", toString(BM)}, |
| 480 | {"{{BK}}", toString(BK)}, |
| 481 | {"{{BN}}", toString(BN)}, |
| 482 | {"{{TM}}", toString(TM)}, |
| 483 | {"{{TN}}", toString(TN)}, |
| 484 | {"{{NUM_TILEA}}", toString(BM * BK / num_threads)}, |
| 485 | {"{{NUM_TILEB}}", toString(BN * BK / num_threads)}, |
| 486 | {"{{TN4}}", toString(TN / 4)}, |
| 487 | {"{{N4}}", toString(N / 4)}, |
| 488 | {"{{BN4}}", toString(BN / 4)}, |
| 489 | }); |
| 490 | if (unrolling) { |
| 491 | std::string unrolledCode = loopUnrolling(codeString); |
| 492 | // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str()); |
| 493 | return {unrolledCode, workgroupSize, precision}; |
| 494 | } else { |
| 495 | return {codeString, workgroupSize, precision}; |
| 496 | } |
| 497 | } |
| 498 | |
| 499 | /* 2D block-tiling with transpose |
| 500 | * |
no test coverage detected