matmul_op.h

Go to the documentation of this file.

#pragma once

#include <string>
#include <vector>

#include "tc/c2/tc_op.h"
#include "tc/library/matmul.h"

namespace caffe2 {

template <typename T, class Context, class Engine = caffe2::DefaultEngine>
class TcMatMulOp : public TcOp<T, Context, Engine> {
 public:
  static constexpr auto description = tc::TC_MATMUL;

  TcMatMulOp(const caffe2::OperatorDef& operator_def, caffe2::Workspace* ws)
      : TcOp<T, Context, Engine>(operator_def, ws) {
    bool trans_a = OperatorBase::GetSingleArgument<int>("trans_a", 0);
    bool trans_b = OperatorBase::GetSingleArgument<int>("trans_b", 0);
    this->tc_ = tc::makeMatmulTc(trans_a, trans_b);
    this->tcName_ = tc::TC_MATMUL_NAME;
  }

  ~TcMatMulOp() override {}

 protected:
  void setupNaiveMappingOptions() override {
    this->mappingOptions_.tile({16, 16, 32})
        .mapToThreads(4, 32)
        .mapToBlocks(32, 32, 32)
        .unroll(1);
  }
};
} // namespace caffe2