matmul.cc

#include "operators/matmul.h"
#include "utils/operator_utils.h"
namespace infini
{

    MatmulObj::MatmulObj(GraphObj *graph, Tensor A, Tensor B, Tensor C, bool transA,
                         bool transB)
        : OperatorObj(OpType::MatMul, TensorVec{A, B}, {C}),
          transA(transA), transB(transB)
    {
        IT_ASSERT(checkValid(graph));
    }

    string MatmulObj::toString() const
    {
        std::ostringstream os;
        os << "Matmul([" << (transA ? "A^T" : "A") << "," << (transB ? "B^T" : "B]")
           << ",A=" << inputs[0]->getGuid()
           << ",B=" << inputs[1]->getGuid() << ",C=" << outputs[0]->getGuid()
           << ",mnk=[" << m << "," << n << "," << k << "])";
        return os.str();
    }

    optional<vector<Shape>> MatmulObj::inferShape(const TensorVec &inputs)
    {
        // =================================== 作业 ===================================
        // TODO：返回经过 matmul 操作后的 shape
        // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm
        // =================================== 作业 ===================================
        //先拿到两个维度
        //这里是多维的矩阵乘法，以最后两维度作为相乘的维度
        const auto A = inputs[0];
        const auto B = inputs[1];
        auto shapeA = A->getDims();
        auto shapeB = B->getDims();

        int rankA = shapeA.size();
        int rankB = shapeB.size();
        //检查是否专职
        if (this->transA && rankA >= 2) {
            std::swap(shapeA[rankA - 1], shapeA[rankA - 2]);
        }
        if (this->transB && rankB >= 2) {
            std::swap(shapeB[rankB - 1], shapeB[rankB - 2]);
        }
        //一般是M,N * N,K = M,K
        int M = shapeA[rankA - 2];
        int K_A = shapeA[rankA - 1];
        int K_B = shapeB[rankB - 2];
        int N = shapeB[rankB - 1];
        //判断形状对不对
        IT_ASSERT(K_A == K_B);
        //其他维度需要广播
        Shape batchA(shapeA.begin(), shapeA.end() - 2);
        Shape batchB(shapeB.begin(), shapeB.end() - 2);

        Shape batchOut = infer_broadcast(batchA, batchB);

        Shape outputShape = batchOut;
        outputShape.push_back(M);
        outputShape.push_back(N);

        return vector<Shape>{outputShape};
    }

} // namespace infini