tuckerconv.cc

TEST_F(ATenCompilationUnitTest, TuckerConv) {
  at::Tensor core = at::CUDA(at::kFloat).rand({32, 10, 10, 10});
  at::Tensor X = at::CUDA(at::kFloat).rand({32, 10, 128});
  at::Tensor Y = at::CUDA(at::kFloat).rand({32, 10, 128});
  at::Tensor Z = at::CUDA(at::kFloat).rand({32, 10, 64});
  at::Tensor weight = at::CUDA(at::kFloat).rand({128, 64, 3, 3});
  std::vector<at::Tensor> inputs = {core, X, Y, Z, weight};
  std::vector<at::Tensor> outputs;

  static constexpr auto TC = R"TC(
    def tuckerconv(float(B,R,R,R) CORE, float(B,R,H) X, float(B,R,W) Y, float(B,R,I) Z, float(O,I,KH,KW) WEIGHT) -> (Out) {
      Out(b,o,h,w) +=! CORE(b,r1,r2,r3) * X(b,r1,h+kh) * Y(b,r2,w+kw) * Z(b,r3,i) * WEIGHT(o,i,kh,kw)
    }
  )TC";

  auto options = tc::MappingOptions::makeNaiveMappingOptions();
  auto name = "tuckerconv";
  std::string cacheFilename = "";
  autotune(cacheFilename, TC, name, inputs, options, {options});
}