matmul.h

Go to the documentation of this file.

#pragma once

#include "tc/library/common.h"

namespace tc {

static constexpr auto TC_MATMUL_NAME = "matmul";

namespace {
static constexpr auto TC_MATMUL = R"TC(
  def matmul(float(${szA0}, ${szA1}) A, float(${szB0}, ${szB1}) B) -> (O) {
    O(i, j) +=! A(${itA0}, ${itA1}) * B(${itB0}, ${itB1})
  }
)TC";
} // namespace

std::string makeMatmulTc(
    bool transposeFirst = false,
    bool transposeSecond = false) {
  std::string tc(TC_MATMUL);
  tc = replaceString(tc, "${szA0}", (transposeFirst ? "K" : "N"));
  tc = replaceString(tc, "${szA1}", (transposeFirst ? "N" : "K"));
  tc = replaceString(tc, "${szB0}", (transposeSecond ? "M" : "K"));
  tc = replaceString(tc, "${szB1}", (transposeSecond ? "K" : "M"));
  tc = replaceString(tc, "${itA0}", (transposeFirst ? "k" : "i"));
  tc = replaceString(tc, "${itA1}", (transposeFirst ? "i" : "k"));
  tc = replaceString(tc, "${itB0}", (transposeSecond ? "j" : "k"));
  tc = replaceString(tc, "${itB1}", (transposeSecond ? "k" : "j"));
  return tc;
}
} // namespace tc