chore: import upstream snapshot with attribution

2026-07-13 13:33:03 +08:00
commit 5b57521aa1
8226 changed files with 3425766 additions and 0 deletions
@@ -0,0 +1,71 @@
+//
+//  TRTMatMul.cpp
+//  MNN
+//
+//  Created by MNN on 2019/09/11.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "TRTMatMul.hpp"
+#include <core/TensorUtils.hpp>
+#include "TRTBackend.hpp"
+
+using namespace std;
+
+namespace MNN {
+
+inline nvinfer1::Dims squeeze_trailing_dims(nvinfer1::Dims const &dims) {
+    nvinfer1::Dims new_dims = dims;
+
+    // for(int i=0; i<new_dims.nbDims; i++)
+    //   printf("kk %d ", new_dims.d[i]);
+    // printf("\n");
+
+    // Note: TRT requires at least one dimension, so we don't squeeze [1]->[]
+    while (new_dims.nbDims > 1 && new_dims.d[new_dims.nbDims - 1] == 1) {
+        // printf("%d ", new_dims.d[new_dims.nbDims-1]);
+        --new_dims.nbDims;
+    }
+    return new_dims;
+}
+
+nvinfer1::MatrixOperation transpose_format(nvinfer1::ITensor *x, bool transpose) {
+    // printf("%d\n", x->getDimensions().nbDims);
+    return transpose ? nvinfer1::MatrixOperation::kTRANSPOSE : nvinfer1::MatrixOperation::kNONE;
+}
+
+TRTMatMul::TRTMatMul(Backend *b, const Op *op, const std::vector<Tensor *> &inputs,
+                     const std::vector<Tensor *> &outputs)
+    : MNN::TRTCommonExecution(b, op) {
+#ifdef TRT_LOG
+    printf("TRTMatMul in\n");
+#endif
+}
+
+std::vector<ITensor *> TRTMatMul::onEncode(const std::vector<ITensor *> &xOp) {
+#ifdef TRT_LOG
+    printf("TRTMatMul in\n");
+#endif
+    auto param       = mOp->main_as_MatMul();
+    auto transpose_a = transpose_format(xOp[0], param->transposeA());
+    auto transpose_b = transpose_format(xOp[1], param->transposeB());
+
+    auto matmul_layer = mTrtBackend->getNetwork()->addMatrixMultiply(*xOp[0], transpose_a, *xOp[1], transpose_b);
+    if (xOp.size() == 2) {
+        return {matmul_layer->getOutput(0)};
+    }
+    auto C = matmul_layer->getOutput(0);
+    auto shuffle =  mTrtBackend->getNetwork()->addShuffle(*(xOp[2]));
+    auto dimReshape = xOp[0]->getDimensions();
+    dimReshape.nbDims = 2;
+    dimReshape.d[0] = 1;
+    dimReshape.d[1] = mInputs[2]->elementSize();    
+    shuffle->setReshapeDimensions(dimReshape);
+    auto biasReshape = shuffle->getOutput(0);
+    auto biasAdd = mTrtBackend->getNetwork()->addElementWise(*C, *biasReshape, ElementWiseOperation::kSUM);
+    return {biasAdd->getOutput(0)};
+}
+
+TRTCreatorRegister<TypedCreator<TRTMatMul>> __matmul_op(OpType_MatMul);
+
+} // namespace MNN