Function createMatmulWithTranspose

examples/matmul/run.cpp:581–614

)";

inline KernelCode createMatmulWithTranspose(const char *shaderTemplate, const size_t M,
                                                const size_t K, const size_t N, const size_t BM,
                                                const size_t BK, const size_t BN,
                                                const size_t TM, const size_t TN,
                                                const Shape &workgroupSize = {256, 1, 1},
                                                NumType precision = kf32) {
  assert(BM % TM == 0);
  assert(BN % TN == 0);
  assert(K % BK == 0);
  assert(M % BM == 0);
  assert(N % BN == 0);
  // # threads = tile A size == tile B size == # threads for computing C
  int num_threads = BM * BN / (TM * TN);
  std::string codeString(shaderTemplate);
  replaceAll(codeString, {{"{{workgroupSize}}", toString(workgroupSize)},
                          {"{{precision}}", toString(precision)},
                          {"{{M}}", toString(M)},
                          {"{{K}}", toString(K)},
                          {"{{N}}", toString(N)},
                          {"{{BM}}", toString(BM)},
                          {"{{BK}}", toString(BK)},
                          {"{{BN}}", toString(BN)},
                          {"{{TM}}", toString(TM)},
                          {"{{TN}}", toString(TN)},
                          {"{{NUM_TILEA}}", toString(BM * BK / num_threads)},
                          {"{{NUM_TILEB}}", toString(BN * BK / num_threads)},
                          {"{{TN4}}", toString(TN / 4)},
                          {"{{N4}}", toString(N / 4)},
                          {"{{BN4}}", toString(BN / 4)},
                          });
  std::string unrolledCode = loopUnrolling(codeString);
  // LOG(kDefLog, kInfo, "Unrolled code:\n%s", unrolledCode.c_str());
  return {unrolledCode, workgroupSize, precision};
}

/**
 * @brief No-Op shader with matmul bindings for performance testing