megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

/**
 * \file dnn/src/common/batched_matrix_mul.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#include "megdnn/oprs.h"
#include "src/common/utils.h"

namespace megdnn {

void BatchedMatrixMulForward::deduce_dtype(DType A, DType B, DType& C) {
    DType C_candi, C_candi2;
    if (A.category() == DTypeCategory::FLOAT) {
        C_candi = A;
    } else if (A.enumv() == DTypeEnum::Int8) {
        C_candi = dtype::Int32();
        C_candi2 = dtype::Int16();
    } else if (A.enumv() == DTypeEnum::QuantizedS8) {
        C_candi = dtype::QuantizedS32(mul_scale(A, B));
    } else if (A.enumv() == DTypeEnum::Quantized8Asymm) {
        C_candi = dtype::QuantizedS32(mul_scale(A, B));
    } else if (A.enumv() == DTypeEnum::Quantized4Asymm) {
        C_candi = dtype::QuantizedS32(mul_scale(A, B));
    }
    if (!C.valid()) {
        C = C_candi;
    }
    megdnn_assert(
            C.valid() && (C == C_candi || C == C_candi2),
            "unsupported BatchedMatMul(%s, %s) -> %s", A.name(), B.name(), C.name());
}
void BatchedMatrixMulForward::deduce_layout(
        const TensorLayout& A, const TensorLayout& B, TensorLayout& C) {
    auto errmsg = [&]() {
        std::string msg;
        msg.append("A=");
        msg.append(A.to_string());
        msg.append(", B=");
        msg.append(B.to_string());
        msg.append(", C=");
        msg.append(C.to_string());
        msg.append(", transposeA=");
        msg.append(std::to_string(m_param.transposeA));
        msg.append(", transposeB=");
        msg.append(std::to_string(m_param.transposeB));
        return msg;
    };
    MEGDNN_MARK_USED_VAR(errmsg);
    auto good_layout = [](const TensorLayout& l) {
        // l.stride[0] == 0 because im2col conv need batched matrixmul and
        // filter tensor need to be broadcasted. It's only implemented in
        // opencl.
        return l.ndim == 3 && l.stride[2] == 1 &&
               l.stride[1] >= static_cast<ptrdiff_t>(l.shape[2]) &&
               (l.shape[0] == 1 ||
                l.stride[0] >= static_cast<ptrdiff_t>(l.shape[1]) * l.stride[1] ||
                l.stride[0] == 0);
    };
    size_t A0, A1, B0, B1;
    A0 = A.shape[1];
    A1 = A.shape[2];
    B0 = B.shape[1];
    B1 = B.shape[2];
    if (m_param.transposeA)
        std::swap(A0, A1);
    if (m_param.transposeB)
        std::swap(B0, B1);
    deduce_dtype(A.dtype, B.dtype, C.dtype);
    megdnn_assert(
            good_layout(A) && good_layout(B) && A1 == B0 && A[0] == B[0] &&
                    A.dtype.enumv() == B.dtype.enumv(),
            "bad input layouts: %s", errmsg().c_str());
    C = TensorLayout(TensorShape({A[0], A0, B1}), C.dtype);
}

void BatchedMatrixMulForward::check_exec(
        const TensorLayout& A, const TensorLayout& B, const TensorLayout& C,
        size_t workspace_in_bytes) {
    TensorLayout C_expect;
    deduce_layout(A, B, C_expect);
    megdnn_assert(
            C_expect.eq_layout(C), "bad layout for C: expect=%s got=%s",
            C_expect.to_string().c_str(), C.to_string().c_str());
    auto required_workspace_in_bytes = get_workspace_in_bytes(A, B, C);
    megdnn_assert(
            workspace_in_bytes >= required_workspace_in_bytes,
            "needed workspace: %zu; got: %zu", required_workspace_in_bytes,
            workspace_in_bytes);
}
}  // namespace megdnn

// vim: syntax=cpp.doxygen